Update app.py

app.py

@@ -32,7 +32,6 @@
 #     else:
 #         st.warning("Please enter a query!")
 
-
 import os
 import streamlit as st
 import torch
@@ -48,11 +47,14 @@ if not hf_token:
 # Check if CUDA is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the model and tokenizer with token authentication
+# Load model in 8-bit mode (less RAM usage)
 MODEL_NAME = "google/gemma-2b-it"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, token=hf_token, torch_dtype=torch.float16 if device == "cuda" else torch.float32, device_map="auto"
+    MODEL_NAME, 
+    token=hf_token,
+    device_map="auto",
+    load_in_8bit=True  # 👈 This reduces RAM usage!
 )
 
 # Streamlit UI
@@ -63,10 +65,10 @@ if st.button("Generate Code"):
     if user_input:
         with st.spinner("⏳ Generating response... Please wait!"):
             inputs = tokenizer(user_input, return_tensors="pt").to(device)
-            output = model.generate(**inputs, max_new_tokens=100)
+            output = model.generate(**inputs, max_new_tokens=50)
             response = tokenizer.decode(output[0], skip_special_tokens=True)
         
         st.subheader("📝 Generated Code:")
-        st.code(response, language="python")  # Display code properly
+        st.code(response, language="python")
     else:
         st.warning("⚠️ Please enter a query!")