File size: 2,023 Bytes
7c39f2c e6f4fec 7c39f2c |
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 |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
# 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",
},
}
# Step 1: Generate context dynamically from schema
def generate_context(schema):
context_lines = []
for table, details in schema.items():
# List table columns
columns = ", ".join(details["columns"])
context_lines.append(f"The {table} table has the following columns: {columns}.")
# Add relationships if present
if details["relations"]:
context_lines.append(f"The {table} table has the following relationship: {details['relations']}.")
return "\n".join(context_lines)
# Generate schema context
schema_context = generate_context(schema)
# Step 2: Load the T5-base-text-to-sql model
model_name = "suriya7/t5-base-text-to-sql" # A model fine-tuned for SQL generation
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# Step 3: Define a natural language query
user_query = "List all orders where the product price is greater than 50."
# Prepare the input for the model
# Adjust the prompt to focus on SQL generation
input_text = f"Convert the following question into an SQL query:\nSchema:\n{schema_context}\n\nQuestion:\n{user_query}"
inputs = tokenizer.encode(input_text, return_tensors="pt")
# Step 4: Generate SQL query
outputs = model.generate(inputs, max_length=128, num_beams=4, early_stopping=True)
generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Step 5: Display the result
print("User Query:", user_query)
print("Generated SQL Query:", generated_sql)
|