Update app.py

app.py

@@ -1,9 +1,11 @@
 # app.py
 import subprocess
+
+# Install dependencies from requirements.txt
 subprocess.run(["pip", "install", "-r", "requirements.txt"])
-import streamlit as st
-from transformers import AutoModelForSeq2SeqLM
 
+import streamlit as st
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 def main():
     st.title("Hugging Face SQL Generator")
@@ -20,17 +22,17 @@ def main():
         st.code(sql_result, language="sql")
 
 def generate_sql(prompt):
-    # Load the "NumbersStation/nsql-350M" model
-    model_name = "NumbersStation/nsql-350M"
+    # Load the "facebook/bart-large-cnn" model
+    model_name = "facebook/bart-large-cnn"
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     # Tokenize and generate SQL
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model(**inputs)
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
+    outputs = model.generate(**inputs)
 
     # Decode the generated SQL
-    sql_query = tokenizer.batch_decode(outputs["output_ids"], skip_special_tokens=True)[0]
+    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     return sql_query