Update app.py

app.py

@@ -24,17 +24,15 @@ def generate_sql(nl_query, schema):
 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
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device, dtype=torch.long)  # Ensure Long dtype
 
     with torch.no_grad():
         output_ids = model.generate(input_ids, max_length=128, pad_token_id=tokenizer.eos_token_id)
 
+    # If model outputs in float16 or bfloat16, convert back to long/int
+    if output_ids.dtype in [torch.float16, torch.bfloat16]:
+        output_ids = output_ids.to(dtype=torch.long)
+
     # Decode and clean the output
     output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     sql_start = output_text.find("SQL Query:") + len("SQL Query:")
@@ -75,7 +73,7 @@ def get_schema(db_path):
     return schema
 
 # --- Streamlit UI ---
-st.title("🔹 AI-Powered Text-to-SQL Generator")
+st.title("AI-Powered Text-to-SQL Generator")
 st.write("Convert natural language questions into SQL queries and execute them.")
 
 # Database selection
@@ -107,7 +105,7 @@ elif db_option == "Enter schema manually":
     schema = st.text_area("Enter your schema:")
 
 if schema:
-    st.subheader("📌 Extracted/Provided Schema:")
+    st.subheader("Extracted/Provided Schema:")
     st.code(schema, language="sql")
 
 # Query input
@@ -118,7 +116,7 @@ if st.button("Generate SQL Query"):
         st.error("❌ Please provide a database or schema first.")
     else:
         sql_query = generate_sql(user_query, schema)
-        st.subheader("🔹 Generated SQL Query:")
+        st.subheader("Generated SQL Query:")
         st.code(sql_query, language="sql")
 
         # Execute SQL if database exists
@@ -137,4 +135,4 @@ if st.button("Generate SQL Query"):
             else:
                 st.error(f"❌ SQL Execution Error: {result}")
         else:
-            st.info("📌 No database provided, only SQL query was generated.")
+            st.info("No database provided, only SQL query was generated.")