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

# Inicialize o tokenizer e o modelo
tokenizer = T5Tokenizer.from_pretrained('t5-small')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = T5ForConditionalGeneration.from_pretrained('cssupport/t5-small-awesome-text-to-sql')
model = model.to(device)
model.eval()

# Função para traduzir português para inglês
def traduzir_para_ingles(texto):
    try:
        traducao = GoogleTranslator(source='pt', target='en').translate(texto)
        return traducao
    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
    inputs = tokenizer(input_prompt, padding=True, truncation=True, return_tensors="pt").to(device)
    
    # Gere a saída
    with torch.no_grad():
        outputs = model.generate(**inputs, max_length=512)
    
    # Decodifique a saída para texto (SQL)
    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return generated_sql

# 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_ingles}\nquery for: {input_prompt_ingles}"
    
    # Gera o SQL
    sql_query = generate_sql(full_prompt)
    return sql_query

# Cria a interface
interface = gr.Interface(
    fn=gerar_sql_interface,
    inputs="text",
    outputs="text",
    title="Gerador de SQL",
    description="Digite uma consulta em linguagem natural (português) e gere a consulta SQL correspondente."
)

# Inicia a interface
interface.launch()