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

# 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 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):
    # Adiciona o prefixo "tables:" e "query for:" automaticamente
    full_prompt = f"tables:\n{input_prompt}\nquery for: {input_prompt}"
    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 e gere a consulta SQL correspondente."
)

# Inicia a interface
interface.launch()