Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
from datasets import load_dataset | |
import random | |
# Load the Spider dataset | |
spider_dataset = load_dataset("spider", split='train') | |
# Load tokenizer and model | |
tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL") | |
model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL") | |
def generate_sql(query): | |
input_text = "translate English to SQL: " + query | |
inputs = tokenizer(input_text, return_tensors="pt", padding=True) | |
outputs = model.generate(**inputs, max_length=512) | |
sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
return sql_query | |
def get_random_example(): | |
example = random.choice(spider_dataset) | |
return example['question'], generate_sql(example['question']) | |
# Create a Gradio interface | |
interface = gr.Interface( | |
fn=generate_sql, | |
inputs=gr.Textbox(lines=2, placeholder="Enter your natural language query here..."), | |
outputs="text", | |
title="NL to SQL with T5", | |
description="This model converts natural language queries into SQL. Enter your query or get a random example from the Spider dataset!" | |
) | |
# Add a button to get a random example from the dataset | |
interface.add_component( | |
gr.Button("Get Random Example"), | |
fn=get_random_example, | |
inputs=[], | |
outputs=[gr.Textbox(label="Random Question"), gr.Textbox(label="Generated SQL")] | |
) | |
# Launch the app | |
if __name__ == "__main__": | |
interface.launch() | |