File size: 2,037 Bytes
4c9f528 |
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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import nltk
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
# Download necessary NLTK data
nltk.download('punkt')
# Load the OmniParser model and tokenizer
model_name = "microsoft/OmniParser"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# SQLBot function to generate SQL commands
def generate_sql_with_omnparser(query):
# Encode the input query
inputs = tokenizer.encode(query, return_tensors="pt")
# Generate SQL command
outputs = model.generate(inputs, max_length=50)
# Decode the generated SQL command
sql_command = tokenizer.decode(outputs[0], skip_special_tokens=True)
return sql_command
# SQLBot function with personality
def sqlbot(query):
gpt_sql_command = generate_sql_with_omnparser(query)
# Adding personality to the bot
if "create table" in query.lower():
response = f"Alright, rolling up my sleeves to create that table for you! Here it is:\n{gpt_sql_command}"
elif "select" in query.lower():
response = f"Got it! Fetching the data you need:\n{gpt_sql_command}"
elif "show tables" in query.lower():
response = f"Let me show you all the tables you've got:\n{gpt_sql_command}"
elif "insert" in query.lower():
response = f"Great! Adding new records as requested:\n{gpt_sql_command}"
elif "update" in query.lower():
response = f"Time to make some updates! Here you go:\n{gpt_sql_command}"
elif "delete" in query.lower():
response = f"Okay, we're deleting those records:\n{gpt_sql_command}"
else:
response = f"Here's what I found:\n{gpt_sql_command}"
return response
# Example usage
user_query = "Create table employees with name age department"
generated_sql = sqlbot(user_query)
print(generated_sql)
# Another example usage
user_query = "Insert into users (name, age) values ('Alice', 30)"
generated_sql = sqlbot(user_query)
print(generated_sql)
|