Update app.py

app.py

@@ -1,21 +1,28 @@
+# app.py
+
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 from deep_translator import GoogleTranslator
 import torch
 
-# Optimización para entornos limitados
 torch.set_num_threads(1)
 torch.set_num_interop_threads(1)
 
-# Modelo público sin token
-tokenizer = T5Tokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
-model = T5ForConditionalGeneration.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
+model_name = "mrm8488/t5-base-finetuned-wikiSQL"
+tokenizer = T5Tokenizer.from_pretrained(model_name)
+model = T5ForConditionalGeneration.from_pretrained(model_name)
+
+# Tu esquema personalizado
+SCHEMA = """
+Tables:
+bodegas(Id, Nombre, Encargado, Telefono, Email, Direccion, Horario, Regional, Latitud, Longitud)
+maestra(CodigoSap, Descripcion, Grupo, Agrupador, Marca, Parte, Operacion, Componente)
+"""
 
-# Función para generar SQL
 def generar_sql(pregunta_espanol):
     try:
         pregunta_ingles = GoogleTranslator(source="es", target="en").translate(pregunta_espanol)
-        prompt = f"translate English to SQL: {pregunta_ingles}"
+        prompt = f"{SCHEMA}\ntranslate English to SQL: {pregunta_ingles}"
         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)
@@ -23,13 +30,12 @@ def generar_sql(pregunta_espanol):
     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 (con T5 de Hugging Face)",
-    description="Convierte preguntas en español a SQL con un modelo entrenado en WikiSQL."
+    title="Texto a SQL con esquema personalizado",
+    description="Escribe una pregunta en español y genera SQL sobre las tablas `bodegas` y `maestra`."
 )
 
 iface.launch()