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

# Configurazione dell'assistente
class ItalianAssistant:
    def __init__(self):
        self.setup_model()
    
    def setup_model(self):
        """Carica il modello per le risposte in italiano"""
        try:
            # Prova a caricare un modello che funziona bene in italiano
            self.chatbot = pipeline(
                "text-generation",
                model="microsoft/DialoGPT-medium",
                torch_dtype=torch.float16,
                device_map="auto"
            )
        except Exception as e:
            print(f"Errore nel caricamento del modello: {e}")
            # Fallback a un modello più semplice
            self.chatbot = pipeline("text-generation", model="distilgpt2")
    
    def generate_response(self, message, history):
        """Genera una risposta in italiano"""
        try:
            # Prompt per migliorare le risposte in italiano
            prompt = f"Sei un assistente AI utile che parla italiano. Rispondi in modo chiaro e preciso.\n\nUtente: {message}\nAssistente:"
            
            response = self.chatbot(
                prompt,
                max_new_tokens=300,
                temperature=0.7,
                do_sample=True,
                pad_token_id=50256,
                repetition_penalty=1.1
            )
            
            # Estrae solo la parte della risposta
            full_text = response[0]['generated_text']
            assistant_response = full_text.split("Assistente:")[-1].strip()
            
            return assistant_response
            
        except Exception as e:
            return f"Mi dispiace, ho riscontrato un errore: {str(e)}"

# Inizializza l'assistente
assistant = ItalianAssistant()

# Crea l'interfaccia Gradio
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🤖 Italian Assistant
        **Il tuo assistente AI personale in italiano!**
        
        Chiedimi qualsiasi cosa: spiegazioni, aiuto con compiti, curiosità, ecc.
        """
    )
    
    chatbot = gr.ChatInterface(
        fn=assistant.generate_response,
        type="messages",
        examples=[
            "Ciao! Chi sei?",
            "Spiegami la teoria della relatività",
            "Aiutami a scrivere una email formale",
            "Quali sono i pianeti del sistema solare?"
        ],
        title="Italian Assistant"
    )

if __name__ == "__main__":
    demo.launch(debug=True, share=True)