File size: 2,036 Bytes
4c735b8 9c8236d c1abd5c 9c8236d 4c735b8 9c8236d 83ce0d2 9c8236d 83ce0d2 9c8236d c1abd5c 9c8236d 4c735b8 9c8236d c8e1797 9c8236d 4c735b8 9c8236d |
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 56 57 58 59 60 61 |
import json
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Define the schema for the database
db_schema = {
"products": ["product_id", "name", "price", "description", "type"],
"orders": ["order_id", "product_id", "quantity", "order_date"],
"customers": ["customer_id", "name", "email", "phone_number"]
}
# Load the model and tokenizer
model_name = "EleutherAI/gpt-neo-2.7B" # You can also use "Llama-2-7b" or another model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
def generate_sql_query(context, question):
"""
This is the description of the database which is given to you, a user can ask
anything related to this database
Args:
context (str): Description of the database schema or table relationships.
question (str): User's natural language query.
Returns:
str: An answer to the question.
"""
# Prepare the prompt
prompt = f"""
Context: {context}
Question: {question}
Query:
"""
# Tokenize input
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to("cuda" if torch.cuda.is_available() else "cpu")
print("Prompt Sent to Model:")
print(prompt)
# Generate SQL query
output = model.generate(inputs.input_ids, max_length=512, num_beams=5, early_stopping=True)
query = tokenizer.decode(output[0], skip_special_tokens=True)
# Extract query from the output
sql_query = query.split("Query:")[-1].strip()
return sql_query
# Schema as a context for the model
schema_description = json.dumps(db_schema, indent=4)
# Example interactive questions
questions = [
"describe the product table for me, what kind of data it is storing and all"
]
for user_question in questions:
print(f"Question: {user_question}")
sql_query = generate_sql_query(schema_description, user_question)
print(f"Generated SQL Query:\n{sql_query}\n") |