HusnaManakkot commited on
Commit
5665aa8
Β·
verified Β·
1 Parent(s): d82d943

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -31
app.py CHANGED
@@ -1,37 +1,21 @@
1
- import gradio as gr
2
- from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
3
- from datasets import load_dataset
4
 
5
- # Load tokenizer and model
6
- tokenizer = AutoTokenizer.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
7
- model = AutoModelForSeq2SeqLM.from_pretrained("hrshtsharma2012/NL2SQL-Picard-final")
8
 
9
- # Initialize the pipeline
10
- nl2sql_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
11
 
12
- # Load a part of the WikiSQL dataset
13
- wikisql_dataset = load_dataset("wikisql", split='train[:5]')
14
 
15
- def generate_sql(query):
16
- results = nl2sql_pipeline(query)
17
- sql_query = results[0]['generated_text']
18
- # Post-process the output to ensure it's a valid SQL query
19
- sql_query = sql_query.replace('<pad>', '').replace('</s>', '').strip()
20
  return sql_query
21
 
22
- # Use examples from the WikiSQL dataset
23
- example_questions = [(question['question'],) for question in wikisql_dataset]
24
-
25
- # Create a Gradio interface
26
- interface = gr.Interface(
27
- fn=generate_sql,
28
- inputs=gr.Textbox(lines=2, placeholder="Enter your natural language query here..."),
29
- outputs="text",
30
- examples=example_questions,
31
- title="NL to SQL with Picard",
32
- description="This model converts natural language queries into SQL using the WikiSQL dataset. Try one of the example questions or enter your own!"
33
- )
34
-
35
- # Launch the app
36
- if __name__ == "__main__":
37
- interface.launch()
 
1
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 
2
 
3
+ # Load the tokenizer and model
4
+ tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base-multi-summarization-sql-en")
5
+ model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5-base-multi-summarization-sql-en")
6
 
7
+ def nl_to_sql(natural_language_query):
8
+ # Tokenize the input query
9
+ input_ids = tokenizer(natural_language_query, return_tensors="pt").input_ids
10
 
11
+ # Generate the SQL query
12
+ output_ids = model.generate(input_ids, max_length=512)[0]
13
 
14
+ # Decode the generated SQL query
15
+ sql_query = tokenizer.decode(output_ids, skip_special_tokens=True)
 
 
 
16
  return sql_query
17
 
18
+ # Example usage
19
+ natural_language_query = "What is the average salary of employees?"
20
+ sql_query = nl_to_sql(natural_language_query)
21
+ print(f"SQL Query: {sql_query}")