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

# ==========================
# 1. Carica il modello italiano
# ==========================
MODEL_NAME = "Mattimax/DACMini-IT"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)

# ==========================
# 2. Lista di argomenti validi
# ==========================
# Tutto ciò che è inerente al negozio (es. prodotti, orari, contatti, marchi)
VALID_KEYWORDS = [
    "intimo", "calze", "collant", "pigiami", "homewear",
    "merceria", "bottoni", "fili", "nastri", "orari",
    "telefono", "email", "marchi", "logo", "contatti"
]

# ==========================
# 3. Funzione di risposta con filtro
# ==========================
def answer_question(user_input):
    # verifica se la domanda contiene almeno una keyword valida
    if not any(keyword.lower() in user_input.lower() for keyword in VALID_KEYWORDS):
        return "Non posso risponderti a questa domanda."

    # prepara input per il modello
    inputs = tokenizer(user_input, return_tensors="pt").to(device)

    # genera la risposta
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=100,
            do_sample=True,
            top_p=0.9,
            temperature=0.7
        )

    # decodifica la risposta
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# ==========================
# 4. Interfaccia Gradio
# ==========================
iface = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(lines=2, placeholder="Scrivi la tua domanda sul negozio..."),
    outputs=gr.Textbox(label="Risposta"),
    title="Assistente Merceria Galli",
    description="Chiedimi solo informazioni sui prodotti e servizi della Merceria Galli."
)

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