app.py

import gradio as gr
from huggingface_hub import InferenceClient
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM
import os

title = "Generador y clasificador de coloquialismos cubanos"

article = """
## Modelo:
Se usa el somosnlp-hackathon-2026/patriae-cuban-dialect-model.

## Fine-tuning:
Se utilizó el dataset somosnlp-hackathon-2026/patriae-cuba-literature-dataset para hacer fine-tuning del modelo gia-uh/cecilia-2b-instruct-v1.

## Participantes del equipo Patriae en SomosNLP 2026:
* **Carlos Luis Barnés Infante (https://huggingface.co/blacknoize404)**
* **Yisel Clavel Quintero (https://huggingface.co/clavel)**
* **Dionis Lopez (https://huggingface.co/inoid)**
"""

HF_TOKEN = os.environ.get("HF_TOKEN", None)

# Cargar el modelo para clasificación
classifier = pipeline(
    "text-classification", 
    model='somosnlp-hackathon-2026/patriae-cuban-dialect-model', 
    token=HF_TOKEN 
)

repo_id='somosnlp-hackathon-2026/patriae-cuban-dialect-model' 

try:
    # Cargar tokenizador y modelo para generación
    tokenizer_chat = AutoTokenizer.from_pretrained(repo_id, token=HF_TOKEN)
    model_chat = AutoModelForCausalLM.from_pretrained(
        repo_id, 
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None,
        token=HF_TOKEN
    )
    
    # Crear pipeline de generación de texto
    generator = pipeline(
        "text-generation",
        model=model_chat,
        tokenizer=tokenizer_chat,
        device=0 if torch.cuda.is_available() else -1
    )
    chat_available = True
    print("✅ Modelo de chat cargado correctamente")
    
except Exception as e:
    print(f"⚠️ Error cargando modelo de chat: {e}")
    chat_available = False
    generator = None


def chat_response(message, history):
    """
    Función para generar respuestas en dialecto cubano usando el modelo de chat
    """
    if not chat_available:
        yield "⚠️ El modelo de chat no está disponible. Usa el clasificador por ahora."
        return
    
    if not message or message.strip() == "":
        yield "🤔 ¿Qué quieres que te diga? Dame algo con qué trabajar."
        return
    
    try:
        # Construir el prompt con contexto
        prompt = f"""Eres un cubano conversando con un amigo. Responde con naturalidad usando el dialecto cubano.

Usuario: {message}
Cubano:"""
        
        # Generar respuesta
        result = generator(
            prompt,
            max_new_tokens=150,
            temperature=0.7,
            top_p=0.95,
            do_sample=True,
            pad_token_id=tokenizer_chat.eos_token_id
        )
        
        # Extraer la respuesta generada
        generated_text = result[0]['generated_text']
        # Quitar el prompt de la respuesta
        response = generated_text.replace(prompt, "").strip()
        
        # Si no hay respuesta, dar un mensaje por defecto
        if not response:
            response = "¡Ay, mi hermano! No sé qué decirte, pero pa' que veas que estoy aquí escuchando."
        
        yield response
        
    except Exception as e:
        yield f"❌ Error en el chat: {str(e)}"


def classify_text(text):
    """
    Función para clasificar texto y detectar coloquialismos cubanos
    """
    if not text or text.strip() == "":
        return "⚠️ Por favor, ingresa un texto para clasificar."
    
    try:
        # Obtener resultados del modelo
        results = classifier(text)
        
        # DEBUG: Imprimir para ver la estructura (puedes eliminar después)
        print(f"Results: {results}")
        print(f"Type: {type(results)}")
        print(f"Length: {len(results)}")
        
        # Verificar que results no esté vacío
        if not results or len(results) == 0:
            return "❌ No se obtuvieron resultados del modelo."
        
        # Obtener el resultado con mayor confianza
        top_result = results[0]
        label = top_result['label']
        score = top_result['score']
        
        # Determinar si es coloquial (LABEL_1 = coloquial, LABEL_0 = no coloquial)
        if label == "LABEL_1":
            es_coloquial = True
            emoji = "🇨🇺"
            tipo = "**COLOQUIAL CUBANO**"
            descripcion = "✅ Este texto contiene expresiones típicas del dialecto cubano."
        elif label == "LABEL_0":
            es_coloquial = False
            emoji = "📚"
            tipo = "**NO COLOQUIAL / ESTÁNDAR**"
            descripcion = "ℹ️ Este texto usa lenguaje estándar, sin expresiones coloquiales cubanas."
        else:
            # Si la etiqueta no es LABEL_0 ni LABEL_1, mostrar la etiqueta original
            emoji = "🔍"
            tipo = f"**{label}**"
            descripcion = f"Etiqueta desconocida: {label}"
        
        # Construir la salida formateada
        output = f"""
### {emoji} Resultado: {tipo}

**Confianza:** {score:.2%}

{descripcion}

"""
        return output
    except Exception as e:
        return f"❌ Error en clasificación: {str(e)}\n\nDetalles técnicos: El modelo devolvió: {results if 'results' in locals() else 'No hay resultados'}"

    
# Interfaz de Chat
chat_interface = gr.ChatInterface(
    fn=chat_response,
    title="💬 Chat en Dialecto Cubano",
    description="Conversa con un modelo entrenado en dialecto cubano.",
    examples=[
        "¿Qué bolá? ¿Cómo andas?",
        "Cuéntame qué hiciste hoy.",
        "¿Cómo se dice 'amigo' en español cubano?",
        "Dime un refrán cubano.",
        "¿Qué opinas de la comida cubana?"
    ]
)

# Interfaz de Clasificación
classifier_interface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(
        label="Texto a clasificar",
        placeholder="Escribe una frase en cubano...",
        lines=5
    ),
    outputs=gr.Markdown(label="Resultado"),
    title="Clasificador de coloquialismos cubanos",
    description="Detecta si un texto contiene expresiones coloquiales cubanas.",
    examples=[
        ["Nos marchábamos sin ti, creyendo que se te habían pegado las sábanas."],
        ["De Varadero venía para acá porque en su casa las broncas con el padre eran de ampanga."],
        ["Los pueblos que no se conocen han de darse prisa para conocerse, como quienes van a pelear juntos."],
        ["El cielo está despejado y hace buen tiempo."],
        ["¿Qué bolá? ¿Cómo andas? ¡Tremenda jodedera!"],

    ]
)

# Combinar ambas interfaces en una sola app
with gr.Blocks(title=title) as demo:
    gr.Markdown(f"# {title}")
    gr.Markdown(article) 

    with gr.Tab("🔍 Clasificador"):
        classifier_interface.render()
    
    with gr.Tab("💬 Chat"):
        chat_interface.render()
    
    gr.Markdown(article)

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