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

# Nome del tuo modello su Hugging Face
model_id = "matteoangeloni/EduRaccoon"

# Config quantizzazione per farlo girare anche su GPU T4 di Hugging Face Spaces
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype="float16"
)

# Carica tokenizer e modello (gestisce automaticamente gli shard .safetensors)
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",              # usa GPU se disponibile
    quantization_config=quant_config,
    torch_dtype="auto"
)

# Crea pipeline di text-generation
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

# Funzione per gestire la chat
def respond(message, history):
    # Ricostruisce la conversazione come testo
    conversation = ""
    if history:
        for user_msg, bot_msg in history:
            conversation += f"User: {user_msg}\nAI: {bot_msg}\n"
    conversation += f"User: {message}\nAI:"

    # Genera la risposta
    result = pipe(
        conversation,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.7,
        top_p=0.95,
        pad_token_id=tokenizer.eos_token_id
    )

    # Estrarre solo la parte dopo "AI:"
    output_text = result[0]["generated_text"]
    response = output_text.split("AI:")[-1].strip()

    return response

# Interfaccia Gradio
chatbot = gr.ChatInterface(
    fn=respond,
    type="messages",
    title="EduRaccoon - Educational AI",
    description="Chatta con EduRaccoon, il tuo assistente AI educativo!"
)

if __name__ == "__main__":
    chatbot.launch()