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

# Carrega o modelo e o tokenizer localmente
model_name = "google/gemma-3-1b-it"  # Substitua pelo caminho local se já baixou

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.7)

def generate_response(message, history):
    messages = [
    [
        {
            "role": "system",
            "content": [{"type": "text", "text": "Você é ELIZA, uma terapeuta que responde com empatia e faz perguntas para entender melhor o paciente."},]
        },
        {
            "role": "user",
            "content": [{"type": "text", "text": message},]
        },
    ],
    ]

    response = pipe(messages)

    return response[0][0]['generated_text'][2]['content']


# Cria a interface Gradio
demo = gr.ChatInterface(
    generate_response,
    title="ELIZA Therapy Session",
    description="Compartilhe seus pensamentos e ELIZA irá ajudar você a refletir sobre eles."
)

# Inicia a interface
demo.launch()