import gradio as gr
from datasets import load_dataset

# Carga el modelo
model = gr.load("models/Qwen/Qwen2.5-Coder-32B-Instruct")

# Carga el split de entrenamiento del dataset
train_dataset = load_dataset("yannelli/laravel-11-qa", split="train")

# Inspección opcional para verificar la estructura
print(train_dataset[0])

# Función para procesar las entradas del chatbot
def chatbot(input_text):
    # Reemplaza "content" con la clave correcta según la inspección
    relevant_docs = [
        doc["text"] for doc in train_dataset
        if input_text.lower() in doc["text"].lower()
    ]
    
    # Respuesta del modelo
    model_response = model(input_text)
    
    # Construye la respuesta personalizada
    if relevant_docs:
        additional_info = "\n\n".join(relevant_docs[:3])  # Limita a 3 documentos relevantes
        return f"{model_response}\n\nDocumentación relevante:\n{additional_info}"
    else:
        return model_response

# Lanza la aplicación con Gradio
gr.Interface(fn=chatbot, inputs="text", outputs="text").launch()