import os
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import login

# Configurar el token de Hugging Face
huggingface_token = os.environ.get('reparbot2')  # Mantenemos tu nombre de token original
if huggingface_token is None:
    raise ValueError("El token de Hugging Face no está configurado en las variables de entorno. Por favor, configura reparbot2")

# Iniciar sesión en Hugging Face
login(huggingface_token)

# Configurar el modelo y tokenizer
model_id = "meta-llama/Llama-3.3-70B-Instruct"  # Mantenemos tu modelo original
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# Función para generar respuestas
def respond_to_query(user_input):
    # Preparar el input
    inputs = tokenizer.encode(user_input, return_tensors="pt").to(model.device)
    
    # Generar respuesta
    outputs = model.generate(
        inputs,
        max_new_tokens=256,
        do_sample=True,
        top_p=0.95,
        top_k=50,
        temperature=0.7,
    )
    
    # Decodificar y retornar la respuesta
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Crear la interfaz de Gradio
interface = gr.Interface(
    fn=respond_to_query,
    inputs=gr.Textbox(label="Tu pregunta"),
    outputs=gr.Textbox(label="Respuesta"),
    title="Chatbot con Llama 3.3",
    description="Haz una pregunta y el modelo te responderá"
)

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