import os
import gradio as gr
from transformers import pipeline

# گرفتن توکن از متغیر محیطی
token = os.environ.get("HF_TOKEN")

# استفاده از مدل سازگار با Transformers
pipe = pipeline(
    "text-generation",
    model="google/gemma-2b-it",  # ✅ مدل سبک و سازگار
    token=token
)

def generate_topics(field, major, keywords, audience, level):
    prompt = f"""
۳ موضوع پایان‌نامه در رشته {field} با گرایش {major} پیشنهاد بده که به کلیدواژه‌های "{keywords}" مربوط باشه و جامعه هدف آن "{audience}" باشد. مقطع: {level}.
موضوعات را فارسی بنویس.
"""
    response = pipe(prompt, max_new_tokens=250)[0]['generated_text']

    # حذف متن prompt از ابتدای پاسخ (در صورتی که مدل آن را تکرار کند)
    if response.startswith(prompt.strip()):
        response = response[len(prompt.strip()):].strip()

    # افزودن متن پایانی ثابت
    response += "\n\nبرای مشاوره و راهنمایی تخصصی با گروه مشاوره کاسپین تماس بگیرید:\n02188252497"

    return response

iface = gr.Interface(
    fn=generate_topics,
    inputs=[
        gr.Textbox(label="رشته"),
        gr.Textbox(label="گرایش"),
        gr.Textbox(label="کلیدواژه‌ها"),
        gr.Textbox(label="جامعه هدف"),
        gr.Dropdown(choices=["کارشناسی ارشد", "دکتری"], label="مقطع")
    ],
    outputs="text",
    title="پیشنهادگر هوشمند موضوع پایان‌نامه کاسپین 🎓"
)

iface.launch()