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

def initialize_chat(model_id: str = "WolfInk/GPT-1.5-High", tokenizer_id: str = "microsoft/DialoGPT-medium"):
    """
    1) Intenta cargar el modelo desde Hugging Face Hub usando model_id.
    2) Carga el tokenizer de DialoGPT-medium.
    """
    # --- Tokenizer (usando DialoGPT-medium) ---
    try:
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
        print(f"✅ Tokenizer '{tokenizer_id}' cargado correctamente.")
    except Exception as e:
        print(f"⚠️ Error al cargar el tokenizer desde '{tokenizer_id}': {e}")
        return None, None

    # --- Modelo (usando el modelo de WolfInk/GPT-1.5-High) ---
    try:
        model = AutoModelForCausalLM.from_pretrained(model_id)
        print(f"✅ Model '{model_id}' loaded correctly.")
    except Exception as e:
        print(f"⚠️ Error to load: '{model_id}': {e}")
        return None, None

    model.eval()
    if torch.cuda.is_available():
        model.to("cuda")

    return tokenizer, model

def generate_response(user_input, tokenizer, model):
    """
    Función para generar la respuesta del modelo a partir de la entrada del usuario.
    """
    # Estructura el prompt en formato de conversación
    conversation = [
        {"from": "human", "value": user_input},
        {"from": "gpt", "value": ""}
    ]

    # Convertir la conversación en un formato que el modelo entienda
    context = ""
    for message in conversation:
        context += f"{message['from']}: {message['value']} "

    # Tokenizamos la entrada
    inputs = tokenizer(context, return_tensors="pt").to(model.device)

    with torch.no_grad():
        # Generamos una respuesta del modelo
        outputs = model.generate(
            **inputs,
            max_new_tokens=20,  # Ajustado para modelos más pequeños
            do_sample=True,
            top_p=0.85,  # Menos diversidad para mejorar coherencia
            temperature=0.5,  # Menos aleatorio, más coherente
            repetition_penalty=1.6,  # Reducido para evitar repeticiones excesivas
            pad_token_id=tokenizer.eos_token_id
        )

    # Decodificar la respuesta generada
    generated = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[-1]:],
        skip_special_tokens=True
    ).strip()

    return generated

def chat_interface():
    # Usar el modelo y tokenizer de Hugging Face
    model_id = "WolfInk/GPT-1.5-High"  # Asegúrate de que este modelo esté disponible en Hugging Face
    tokenizer_id = "microsoft/DialoGPT-medium"  # Usar el tokenizer de DialoGPT-medium

    tokenizer, model = initialize_chat(model_id, tokenizer_id)

    if tokenizer and model:
        return gr.Interface(
            fn=lambda user_input: generate_response(user_input, tokenizer, model),
            inputs=gr.Textbox(label="👤 You: ", placeholder="Write your message..."),
            outputs=gr.Textbox(label="🤖 GPT-1.5 Response: "),
            title="GPT-1.5 High Chatbot",
            description="Interact with GPT-1.5 High",
            theme="default"
        ).launch()

    else:
        return "❌ Error."

if __name__ == "__main__":
    chat_interface()