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

# --- مدل متن باز HuggingFace ---
MODEL_NAME = "openai/gpt-oss-20b"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# --- تابع چت ---
def chat_cao(user_input):
    inputs = tokenizer(user_input, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=200)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# --- ظاهر HTML/Frontend ---
css = """
body { background-color: #f5f7fa; font-family: 'Arial', sans-serif; }
h1 { color: #6a5acd; text-align: center; }
.gradio-container { border-radius: 12px; padding: 20px; }
input, textarea { border-radius: 8px; }
"""

iface = gr.Interface(
    fn=chat_cao,
    inputs=gr.Textbox(lines=2, placeholder="پیام خودتو بنویس..."),
    outputs="text",
    title="CAO - هوش مصنوعی بی‌سانسور",
    description="چت فارسی و انگلیسی، آماده اضافه کردن کدنویسی، فایل و تصویر",
    theme="default",
    css=css
)

iface.launch()