import gradio as gr
import spaces
import torch

from transformers import AutoTokenizer, AutoModelForCausalLM


MODEL_ID = "ddfws/Rezaeian-StatsAI"


print("Loading tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID
)


print("Loading model...")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    dtype=torch.float16
)

model.eval()

print("MODEL READY")



def format_answer(text):

    # مرتب سازی خروجی فارسی
    text = text.replace("\n", "<br>")

    # فرمول های LaTeX
    text = text.replace("\\[", "<div class='formula'>")
    text = text.replace("\\]", "</div>")

    text = text.replace("$$", "<div class='formula'>")

    return text



@spaces.GPU
def generate(message, history):

    prompt = ""


    for user, assistant in history:

        prompt += f"User: {user}\nAssistant: {assistant}\n"


    prompt += f"User: {message}\nAssistant:"


    inputs = tokenizer(
        prompt,
        return_tensors="pt"
    )


    inputs = {
        k: v.to(model.device)
        for k, v in inputs.items()
    }


    with torch.no_grad():

        output = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.9,
            do_sample=True
        )


    result = tokenizer.decode(
        output[0],
        skip_special_tokens=True
    )


    answer = result.split("Assistant:")[-1]


    return format_answer(answer)



demo = gr.ChatInterface(
    fn=generate,
    title="Rezaeian StatsAI",
    description="دستیار هوشمند آمار و احتمال مهندسی"
)



demo.launch(
    server_name="0.0.0.0",
    css="""

    .gradio-container {
        direction: rtl;
        max-width: 1200px !important;
    }


    .message {
        direction: rtl !important;
        text-align: right !important;
        font-family: Tahoma, Arial, sans-serif !important;
        font-size:18px !important;
        line-height:2;
    }


    textarea {
        direction: rtl !important;
        text-align:right !important;
        font-size:18px !important;
    }


    .formula {
        direction:ltr !important;
        text-align:center !important;
        font-family: Cambria Math, Times New Roman, serif !important;
        font-size:20px !important;
        padding:12px;
        margin:12px;
    }


    footer {
        display:none !important;
    }

    """
)