Update app.py

app.py

@@ -7,26 +7,31 @@ model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 
 def generate_sql(payload):
-    # Extract components from payload
+    # Extract parts from the JSON payload
     question = payload.get("question", "")
     schema = payload.get("schema", "")
-    
-    # Build model input
-    full_prompt = f"Question: {question} Schema: {schema}"
-    
-    inputs = tokenizer(full_prompt, return_tensors="pt")
+    sample_rows = payload.get("sample_rows", [])
+
+    # Convert sample rows into a single string
+    sample_str = " ".join([str(row) for row in sample_rows]) if sample_rows else ""
+
+    # Build model input prompt
+    prompt = f"Question: {question} Schema: {schema} Sample Rows: {sample_str}"
+
+    # Tokenize and generate
+    inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(**inputs, max_length=512)
     generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
+
     return generated_sql
 
-# Define expected input as a JSON object (dict)
+# Gradio interface
 demo = gr.Interface(
     fn=generate_sql,
-    inputs=gr.JSON(label="Input JSON (with 'question' and 'schema')"),
+    inputs=gr.JSON(label="Input JSON (question, schema, sample_rows)"),
     outputs="text",
     title="Text-to-SQL Generator",
-    description="Input a JSON with your natural language question and database schema. Output is SQL."
+    description="Enter a JSON object with 'question', 'schema', and optional 'sample_rows'. The model will generate SQL."
 )
 
 demo.launch()