sqlgenerator / app.py
JeffersonCorreiax's picture
Update app.py
c25cc78 verified
raw
history blame
987 Bytes
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()