import gradio as gr
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 chat(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"
    ).to(model.device)


    with torch.no_grad():

        result = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.7,
            do_sample=True
        )


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


    return text.split("Assistant:")[-1]



demo = gr.ChatInterface(
    fn=chat,
    title="Rezaeian StatsAI"
)


demo.launch(
    server_name="0.0.0.0",
    server_port=7860
)