updated code

README.md

@@ -5,7 +5,7 @@ colorFrom: red
 colorTo: red
 sdk: gradio
 sdk_version: 5.11.0
-app_file: database.py
+app_file: t5.py
 pinned: false
 python: 3.9
 ---

db.py

@@ -0,0 +1,55 @@
+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 = "mrm8488/t5-base-finetuned-wikiSQL"  # 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)

t5.py

@@ -0,0 +1,21 @@
+from transformers import AutoModelWithLMHead, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
+model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
+
+def get_sql(query):
+  input_text = "translate English to SQL: %s </s>" % query
+  features = tokenizer([input_text], return_tensors='pt')
+
+  output = model.generate(input_ids=features['input_ids'], 
+               attention_mask=features['attention_mask'])
+  
+  return tokenizer.decode(output[0])
+
+query = "How many models were finetuned using BERT as base model?"
+
+res = get_sql(query)
+
+print(res)
+
+# output: 'SELECT COUNT Model fine tuned FROM table WHERE Base model = BERT'