Fixed hallucinations

app.py

@@ -14,7 +14,7 @@ def generate_sql(context: str, query: str) -> str:
     """
     Generates a SQL query given the provided context and natural language query.
     Constructs a prompt from the inputs, then performs deterministic generation
-    with beam search.
+    with beam search and repetition handling.
     """
     prompt = f"""Context:
 {context}
@@ -25,24 +25,28 @@ Query:
 Response:
 """
     # Tokenize the prompt and move to device
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
     
     # Ensure decoder_start_token_id is set for encoder-decoder generation
     if model.config.decoder_start_token_id is None:
         model.config.decoder_start_token_id = tokenizer.pad_token_id
 
-    # Generate the SQL output
+    # Generate the SQL output with optimized parameters
     generated_ids = model.generate(
         input_ids=inputs["input_ids"],
         decoder_start_token_id=model.config.decoder_start_token_id,
-        max_new_tokens=250,
-        temperature=0.0,      # Deterministic output
-        num_beams=3,          # Beam search for improved quality
-        early_stopping=True,  # Stop when output is complete
+        max_new_tokens=100,   
+        temperature=0.1,      
+        num_beams=5,          
+        repetition_penalty=1.2, 
+        early_stopping=True, 
     )
     
-    # Decode and return the generated SQL statement
-    return tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    # Decode and clean the generated SQL statement
+    generated_sql = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    generated_sql = generated_sql.split(";")[0] + ";"  # ✅ Ensures only the first valid SQL query is returned
+
+    return generated_sql
 
 # Create Gradio interface with two input boxes: one for context and one for query
 iface = gr.Interface(
@@ -53,7 +57,9 @@ iface = gr.Interface(
     ],
     outputs="text",
     title="Text-to-SQL Generator",
-    description="Enter your own context (e.g., database schema and sample data) and a natural language query. The model will generate the corresponding SQL statement."
+    description="Enter your own context (e.g., database schema and sample data) and a natural language query. The model will generate the corresponding SQL statement.",
+    theme="compact",
+    allow_flagging="never"
 )
 
 iface.launch()