Update app.py

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from datetime import datetime
 
-# Custom CSS for a Grok/ChatGPT-like look
+# Custom CSS for UI
 st.markdown("""
     <style>
     .main { background-color: #f9f9f9; padding: 20px; }
@@ -57,29 +57,44 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
-# Load model and tokenizer
-try:
+# Cache model and tokenizer to avoid reloading
+@st.cache_resource
+def load_model_and_tokenizer():
     checkpoint = "Salesforce/codegen-350M-mono"
-    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-    model = AutoModelForCausalLM.from_pretrained(checkpoint)
-except Exception as e:
-    st.error(f"Error loading model: {e}")
+    try:
+        st.write("Loading tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+        st.write("Loading model...")
+        model = AutoModelForCausalLM.from_pretrained(checkpoint)
+        st.write("Model and tokenizer loaded successfully!")
+        return tokenizer, model
+    except Exception as e:
+        st.error(f"Failed to load model/tokenizer: {e}")
+        return None, None
+
+# Load model and tokenizer once
+tokenizer, model = load_model_and_tokenizer()
+if tokenizer is None or model is None:
     st.stop()
 
 # Function to generate code
 def generate_code(description):
     prompt = f"Generate Python code for the following task: {description}\n"
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(
-        **inputs,
-        max_length=500,
-        num_return_sequences=1,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    code = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return code[len(prompt):].strip()
+    try:
+        outputs = model.generate(
+            **inputs,
+            max_length=500,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return code[len(prompt):].strip()
+    except Exception as e:
+        st.error(f"Error generating code: {e}")
+        return "Error: Could not generate code."
 
-# Initialize chat history in session state
+# Initialize chat history
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
@@ -95,13 +110,12 @@ with st.container():
         height=150
     )
     
-    col1, col2 = st.columns([1, 1])  # Adjusted for two buttons
+    col1, col2 = st.columns([1, 1])
     with col1:
         if st.button("Generate"):
             if description.strip():
                 with st.spinner("Thinking..."):
                     generated_code = generate_code(description)
-                # Append to chat history with timestamp
                 st.session_state.chat_history.append({
                     "input": description,
                     "output": generated_code,
@@ -121,7 +135,6 @@ with st.container():
         for chat in st.session_state.chat_history:
             st.markdown(f'<div class="chat-message"><strong>You ({chat["time"]}):</strong> {chat["input"]}</div>', unsafe_allow_html=True)
             st.markdown(f'<div class="code-output">{chat["output"]}</div>', unsafe_allow_html=True)
-            st.markdown("---")  # Separator for readability
+            st.markdown("---")
 
-# Optional tip at the bottom
 st.info("Tip: Check the generated code for accuracy before using it!")