Spaces:
Sleeping
Sleeping
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() |