import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch

# Limitar el uso de CPU para servidores lentos
torch.set_num_threads(1)
torch.set_num_interop_threads(1)

# Cargar modelo
tokenizer = T5Tokenizer.from_pretrained("cssupport/t5-small-awesome-text-to-sql")
model = T5ForConditionalGeneration.from_pretrained("cssupport/t5-small-awesome-text-to-sql")

# Esquema de base de datos simplificado
SCHEMA = """
Database schema:
Table bodegas(Id, Nombre, Encargado, Telefono, Email, Direccion, Horario, Regional, Latitud, Longitud)
Table maestra(CodigoSap, Descripcion, Grupo, Agrupador, Marca, Parte, Operacion, Componente)
"""

# Función principal
def generar_sql(pregunta_espanol):
    try:
        # Crear prompt claro y directo
        prompt = f"Esquema de base de datos:\n{SCHEMA}\nPregunta en español: {pregunta_espanol}\nGenera la consulta SQL correspondiente."

        # Generar la consulta SQL
        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
        output = model.generate(input_ids, max_length=128)
        sql = tokenizer.decode(output[0], skip_special_tokens=True)
        
        return sql
    except Exception as e:
        return f"Error: {str(e)}"

# Interfaz Gradio
iface = gr.Interface(
    fn=generar_sql,
    inputs=gr.Textbox(lines=3, label="Pregunta en español"),
    outputs=gr.Textbox(label="Consulta SQL generada"),
    title="Texto a SQL (entrada en español)",
    description="Escribe una pregunta en español sobre la base de datos y obtén la consulta SQL."
)

iface.launch()