File size: 1,360 Bytes
54e8483
599a21f
75e09ad
54e8483
4c23181
eefed22
 
54e8483
4c23181
 
9b65c50
75e09ad
 
 
4c23181
eefed22
 
 
 
 
4c23181
b503163
3b69718
75e09ad
 
fd893e6
4c23181
 
b503163
5db5fa6
4c23181
fd893e6
eefed22
 
3b69718
 
 
2a369c5
eefed22
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from datasets import load_dataset

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")

# Initialize the pipeline
nl2sql_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

# Load a part of the Spider dataset
spider_dataset = load_dataset("spider", split='train[:5]')

def generate_sql(query):
    # Format the input for the model
    input_text = f"translate English to SQL: {query}"
    # Run the pipeline
    results = nl2sql_pipeline(input_text)
    # Extract the SQL query
    sql_query = results[0]['generated_text']
    return sql_query

# Use examples from the Spider dataset
example_questions = [(question['question'],) for question in spider_dataset]

# Create a Gradio interface
interface = gr.Interface(
    fn=generate_sql,
    inputs=gr.Textbox(lines=2, placeholder="Enter your natural language query here..."),
    outputs="text",
    examples=example_questions,
    title="NL to SQL with T5",
    description="This model converts natural language queries into SQL. Try one of the example questions or enter your own!"
)

# Launch the app
if __name__ == "__main__":
    interface.launch()