File size: 2,285 Bytes
f1b2798 |
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 62 63 64 65 66 67 68 |
from transformers import AutoTokenizer, AutoModelForCausalLM
# Load model and tokenizer
model_name = "EleutherAI/gpt-neo-2.7B" # Replace with a suitable model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Example schema
schema = {
"products": {
"columns": ["product_id", "name", "price", "category_id"],
"relations": "category_id -> categories.id",
},
"categories": {
"columns": ["id", "category_name"],
"relations": None,
},
"orders": {
"columns": ["order_id", "customer_name", "product_id", "order_date"],
"relations": "product_id -> products.product_id",
},
}
def generate_context(schema):
"""
Generate context dynamically from the schema.
"""
context = "### Database Schema ###\n\n"
for table, details in schema.items():
context += f"Table: {table}\nColumns: {', '.join(details['columns'])}\n"
if details.get("relations"):
context += f"Relations: {details['relations']}\n"
context += "\n"
context += "### Instructions ###\n"
context += (
"Generate SQL queries based on the user's question. "
"Understand the schema to identify relevant tables and relationships. "
"If the question involves multiple tables, use appropriate joins.\n"
)
return context
# Generate dynamic context
context = generate_context(schema)
def answer_question(context, question):
"""
Generate an SQL query or database-related response using the model.
"""
prompt = f"{context}\n\nUser Question: {question}\nSQL Query or Answer:"
inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
outputs = model.generate(inputs.input_ids, max_length=256, num_beams=5, early_stopping=True)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Interactive loop
print("Database Assistant is ready. Ask your questions!")
# 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}")
response = answer_question(context, user_question)
print("\nGenerated Response:\n", response, "\n")
|