Update app.py

app.py

@@ -33,7 +33,6 @@
 #         st.warning("Please enter a query!")
 
 
-
 import os
 import streamlit as st
 import torch
@@ -52,7 +51,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Load the model and tokenizer with token authentication
 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 = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME, token=hf_token, torch_dtype=torch.float16 if device == "cuda" else torch.float32, device_map="auto"
+)
 
 # Streamlit UI
 st.title("Gemma-2B Code Assistant")
@@ -60,10 +61,12 @@ user_input = st.text_area("Enter your coding query:")
 
 if st.button("Generate Code"):
     if user_input:
-        inputs = tokenizer(user_input, return_tensors="pt").to(device)
-        output = model.generate(**inputs, max_new_tokens=100)
-        response = tokenizer.decode(output[0], skip_special_tokens=True)
-        st.write(response)
+        with st.spinner("⏳ Generating response... Please wait!"):
+            inputs = tokenizer(user_input, return_tensors="pt").to(device)
+            output = model.generate(**inputs, max_new_tokens=100)
+            response = tokenizer.decode(output[0], skip_special_tokens=True)
+        
+        st.subheader("📝 Generated Code:")
+        st.code(response, language="python")  # Display code properly
     else:
-        st.warning("Please enter a query!")
-
+        st.warning("⚠️ Please enter a query!")