# نصب کتابخانه‌ها (در HuggingFace Space خودکار نصب می‌شود)
# !pip install transformers accelerate gradio

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

MODEL_NAME = "openai/gpt-oss-120b"

# لود tokenizer و مدل با کمترین مصرف حافظه
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",        # توزیع خودکار بین CPU/GPU
    torch_dtype=torch.float16,
    offload_folder="offload", # ذخیره موقت روی دیسک
)

def chat_with_bot(user_input, history=[]):
    # اضافه کردن تاریخچه
    history_text = ""
    for user_msg, bot_msg in history:
        history_text += f"User: {user_msg}\nBot: {bot_msg}\n"
    prompt = history_text + f"User: {user_input}\nBot:"

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(
        **inputs, 
        max_new_tokens=300, 
        temperature=0.7, 
        top_p=0.9
    )
    bot_reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
    bot_reply = bot_reply.split("Bot:")[-1].strip()

    history.append((user_input, bot_reply))
    return bot_reply, history

# رابط گرافیکی با Gradio
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox(placeholder="پیام خودت رو وارد کن...")
    clear = gr.Button("پاک کردن تاریخچه")

    def respond(message, chat_history):
        reply, history = chat_with_bot(message, chat_history)
        return history, history

    msg.submit(respond, [msg, chatbot], [chatbot])
    clear.click(lambda: [], None, chatbot)

demo.launch()