from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch
import gradio as gr

# ✅ استفاده از نسخه موجود و عمومی PersianMind
model_name = "universitytehran/PersianMind-v1.0"

# بارگذاری tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)

# بررسی وجود GPU و فعال‌سازی quantization در صورت امکان
if torch.cuda.is_available():
    quant_config = BitsAndBytesConfig(load_in_8bit=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        quantization_config=quant_config,
        device_map="auto"
    )
else:
    model = AutoModelForCausalLM.from_pretrained(model_name)

# تابع استخراج کلمات کلیدی با پرامپت
def extract_keywords(text):
    prompt = f"""متن زیر را بخوان و فقط کلمات کلیدی مهم آن را با کاما جدا کن:\nمتن: {text}\nکلمات کلیدی:"""
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output = model.generate(
        **inputs,
        max_new_tokens=50,
        do_sample=False,
        repetition_penalty=1.2
    )
    keywords = tokenizer.decode(output[0], skip_special_tokens=True)
    if "کلمات کلیدی:" in keywords:
        keywords = keywords.split("کلمات کلیدی:")[-1].strip()
    return keywords if keywords else "کلمه کلیدی‌ای پیدا نشد."

# رابط کاربری Gradio
interface = gr.Interface(
    fn=extract_keywords,
    inputs=gr.Textbox(label="متن خود را وارد کنید", placeholder="مثلاً: هوش مصنوعی آینده تکنولوژی را تغییر می‌دهد."),
    outputs=gr.Textbox(label="کلمات کلیدی"),
    title="استخراج کلمات کلیدی با PersianMind",
    description="این برنامه با استفاده از مدل PersianMind-v1.0 کلمات کلیدی متن فارسی را استخراج می‌کند."
)

interface.launch()