File size: 1,512 Bytes
effc31c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration
from deep_translator import GoogleTranslator
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
SCHEMA = """

Database schema:

Table customers(id, name, age, country)

Table orders(id, customer_id, amount)

"""
# Función principal
def generar_sql(pregunta_espanol):
   try:
       # Traducir pregunta a inglés
       pregunta_ingles = GoogleTranslator(source="es", target="en").translate(pregunta_espanol)
       # Crear prompt
       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)
       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()