import gradio as gr
from transformers import pipeline

# ساخت pipeline تولید متن با مدل GPT-2 فارسی
generator = pipeline("text-generation", model="HooshvareLab/gpt2-fa")

def generate_text(prompt, max_length):
    """
    تابع دریافت متن ورودی (prompt) و تولید متن تکمیلی با حد اکثر طول مشخص شده.
    """
    # فراخوانی مدل برای تولید متن
    generated = generator(prompt, max_length=int(max_length), do_sample=True)
    # مدل خروجی را در قالب یک دیکشنری برمی‌گرداند؛ به همین دلیل اولین خروجی را انتخاب می‌کنیم.
    return generated[0]['generated_text']

# طراحی رابط کاربری با استفاده از Gradio
iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(lines=5, placeholder="متن ورودی خود را وارد کنید...", label="متن ورودی (Prompt)"),
        gr.Slider(minimum=20, maximum=200, step=1, value=100, label="حداکثر طول متن تولیدی")
    ],
    outputs="text",
    title="تولید متن فارسی با HooshvareLab/gpt2-fa",
    description="با وارد کردن یک متن اولیه، مدل GPT-2 فارسی متن تکمیلی تولید می‌کند."
)

if __name__ == "__main__":
    iface.launch()