Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 import gradio as gr
+from googletrans import Translator
 
 # Inicialize o tokenizer e o modelo
 tokenizer = T5Tokenizer.from_pretrained('t5-small')
@@ -9,6 +10,18 @@ model = T5ForConditionalGeneration.from_pretrained('cssupport/t5-small-awesome-t
 model = model.to(device)
 model.eval()
 
+# Inicialize o tradutor
+translator = Translator()
+
+# Função para traduzir português para inglês
+def traduzir_para_ingles(texto):
+    try:
+        traducao = translator.translate(texto, src='pt', dest='en')
+        return traducao.text
+    except Exception as e:
+        print(f"Erro na tradução: {e}")
+        return texto  # Retorna o texto original em caso de erro
+
 # Função para gerar SQL
 def generate_sql(input_prompt):
     # Tokenize a entrada
@@ -25,8 +38,13 @@ def generate_sql(input_prompt):
 
 # Interface Gradio
 def gerar_sql_interface(input_prompt):
+    # Traduz o prompt de português para inglês
+    input_prompt_ingles = traduzir_para_ingles(input_prompt)
+    
     # Adiciona o prefixo "tables:" e "query for:" automaticamente
-    full_prompt = f"tables:\n{input_prompt}\nquery for: {input_prompt}"
+    full_prompt = f"tables:\n{input_prompt_ingles}\nquery for: {input_prompt_ingles}"
+    
+    # Gera o SQL
     sql_query = generate_sql(full_prompt)
     return sql_query
 
@@ -36,7 +54,7 @@ interface = gr.Interface(
     inputs="text",
     outputs="text",
     title="Gerador de SQL",
-    description="Digite uma consulta em linguagem natural e gere a consulta SQL correspondente."
+    description="Digite uma consulta em linguagem natural (português) e gere a consulta SQL correspondente."
 )
 
 # Inicia a interface