Deploy SQLCoder with Streamlit

app.py

@@ -9,9 +9,20 @@ st.set_page_config(page_title="SQLCoder", layout="wide")
 def load_model():
     model_name = "defog/sqlcoder-7b-2"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",  # Automatically uses GPU if available
+        low_cpu_mem_usage=True  # Optimizes memory usage
+    )
     return tokenizer, model
 
+# @st.cache_resource
+# def load_model():
+#     model_name = "defog/sqlcoder-7b-2"
+#     tokenizer = AutoTokenizer.from_pretrained(model_name)
+#     model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+#     return tokenizer, model
+
 tokenizer, model = load_model()
 
 # Define the function to generate code

requirements.txt

@@ -1,3 +1,4 @@
 transformers
 torch
 streamlit
+accelerate