import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import spaces  # Import spaces for ZeroGPU

# Load the model and tokenizer
model_name = "AndreaAlessandrelli4/AvvoChat_AITA_v03"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Function to load the model with GPU allocation
@spaces.GPU
def load_model():
    model = AutoModelForCausalLM.from_pretrained(model_name)
    model.to("cuda")
    return model

model = load_model()

# Function for the chatbot interaction
def chat(input_text):
    inputs = tokenizer.encode(input_text, return_tensors="pt").to("cuda")
    outputs = model.generate(inputs, max_length=100, num_return_sequences=1)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Custom CSS for the interface
custom_css = """
#component-0 input[type="text"] {
    border: 1px solid #ccc;
    border-radius: 4px;
    padding: 8px;
    font-size: 16px;
    width: 100%;
    box-sizing: border-box;
}
#component-0 .output-textbox {
    border: 1px solid #ccc;
    border-radius: 4px;
    padding: 8px;
    font-size: 16px;
    min-height: 150px;
    width: 100%;
    box-sizing: border-box;
}
body {
    font-family: Arial, sans-serif;
    background-color: #f5f5f5;
}
#component-0 {
    max-width: 600px;
    margin: 0 auto;
    padding: 20px;
    background-color: #fff;
    border: 1px solid #ddd;
    border-radius: 4px;
    box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
}
#component-0 h1 {
    font-size: 24px;
    margin-bottom: 20px;
}
#component-0 .footer {
    margin-top: 20px;
    text-align: center;
    color: #888;
}
"""

# Configure the Gradio interface
iface = gr.Interface(
    fn=chat,
    inputs=gr.inputs.Textbox(lines=1, placeholder="Scrivi il tuo messaggio qui..."),
    outputs="text",
    title="AvvoChat AITA",
    description="Chatbot basato sul modello AndreaAlessandrelli4/AvvoChat_AITA_v03",
    css=custom_css,
)

# Launch the interface
iface.launch()