Update app.py

app.py

@@ -25,6 +25,13 @@ SQL Query:
 """
 
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+
+    # Ensure dtype consistency to avoid errors
+    if model.dtype == torch.float16:
+        input_ids = input_ids.half()  # Convert to float16
+    elif model.dtype == torch.bfloat16:
+        input_ids = input_ids.bfloat16()  # Convert to bfloat16
+
     with torch.no_grad():
         output_ids = model.generate(input_ids, max_length=128, pad_token_id=tokenizer.eos_token_id)
 
@@ -37,6 +44,7 @@ SQL Query:
     sql_query = re.sub(r"```sql|```", "", sql_query).split("###")[0].strip()
     return sql_query
 
+
 def execute_sql(sql_query, db_path):
     """Execute the generated SQL query on the provided database."""
     try: