Update app.py

app.py

@@ -5,13 +5,13 @@ import subprocess
 subprocess.run(["pip", "install", "-r", "requirements.txt"])
 
 import streamlit as st
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 def main():
     st.title("Hugging Face SQL Generator")
 
     # Get user input
-    prompt = st.text_area("Enter your prompt:")
+    prompt = st.text_area("Enter your SQL prompt:")
 
     if st.button("Generate SQL"):
         # Call a function to generate SQL using the Hugging Face model
@@ -24,15 +24,13 @@ def main():
 def generate_sql(prompt):
     # Load the "NumbersStation/nsql-350M" model
     model_name = "NumbersStation/nsql-350M"
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     # Tokenize and generate SQL
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
-    outputs = model.generate(**inputs)
-
-    # Decode the generated SQL
-    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    generated_ids = model.generate(input_ids, max_length=500)
+    sql_query = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
 
     return sql_query