Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| import torch | |
| MODEL_NAME = "gaussalgo/T5-LM-Large-text2sql-spider" | |
| print("Loading model...") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME) | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| model = model.to(device) | |
| print(f"Model ready on {device}") | |
| def generate_sql(question: str, context: str) -> str: | |
| input_text = f"{question} | {context}" | |
| inputs = tokenizer( | |
| input_text, | |
| return_tensors='pt', | |
| max_length=512, | |
| truncation=True | |
| ).to(device) | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=128, | |
| num_beams=4, | |
| early_stopping=True | |
| ) | |
| sql = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return sql | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# NL2SQL API") | |
| with gr.Row(): | |
| question = gr.Textbox(label="Question") | |
| context = gr.Textbox(label="Context") | |
| output = gr.Textbox(label="SQL") | |
| btn = gr.Button("Submit") | |
| btn.click(fn=generate_sql, inputs=[question, context], outputs=output) | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |