import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Carregar o modelo e o tokenizador
tokenizer = AutoTokenizer.from_pretrained("defog/sqlcoder")
model = AutoModelForCausalLM.from_pretrained("defog/sqlcoder")

# Função para gerar SQL a partir de linguagem natural
def gerar_sql(consulta_natural):
    # Tokenizar a consulta em linguagem natural
    inputs = tokenizer(consulta_natural, return_tensors="pt", max_length=512, truncation=True)
    
    # Gerar a consulta SQL
    outputs = model.generate(**inputs, max_length=512)
    
    # Decodificar a saída para texto
    sql_gerado = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return sql_gerado

# Interface Gradio
interface = gr.Interface(
    fn=gerar_sql,
    inputs="text",
    outputs="text",
    title="Gerador de SQL com SQLCoder",
    description="Digite uma consulta em linguagem natural e gere a consulta SQL correspondente."
)

# Iniciar a interface
interface.launch()