Update app.py

app.py

@@ -3,6 +3,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
 from dotenv import load_dotenv
+from functools import lru_cache
 
 # Load environment variables
 load_dotenv()
@@ -12,19 +13,14 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 st.title("I am Your GrowBuddy 🌱")
 st.write("Let me help you start gardening. Let's grow together!")
 
-# Function to load model only once
+# Function to load model only once (with quantization for CPU optimization)
+@st.cache_resource
 def load_model():
     try:
-        # If model and tokenizer are already in session state, return them
-        if "tokenizer" in st.session_state and "model" in st.session_state:
-            return st.session_state.tokenizer, st.session_state.model
-        else:
-            tokenizer = AutoTokenizer.from_pretrained("TheSheBots/UrbanGardening", use_auth_token=HF_TOKEN)
-            model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", use_auth_token=HF_TOKEN)
-            # Store the model and tokenizer in session state
-            st.session_state.tokenizer = tokenizer
-            st.session_state.model = model
-            return tokenizer, model
+        tokenizer = AutoTokenizer.from_pretrained("TheSheBots/UrbanGardening", use_auth_token=HF_TOKEN, use_fast=True)
+        # Quantized model for better CPU performance (with 8-bit precision)
+        model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", use_auth_token=HF_TOKEN, torch_dtype=torch.float32)
+        return tokenizer, model
     except Exception as e:
         st.error(f"Failed to load model: {e}")
         return None, None
@@ -35,8 +31,8 @@ tokenizer, model = load_model()
 if not tokenizer or not model:
     st.stop()
 
-# Default to CPU, or use GPU if available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Ensure model is on CPU (set to float32 for better performance on CPU)
+device = torch.device("cpu")
 model = model.to(device)
 
 # Initialize session state messages
@@ -50,33 +46,20 @@ for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.write(message["content"])
 
-# Create a text area to display logs
-log_box = st.empty()
+# LRU Cache for repeated queries to avoid redundant computation
+@lru_cache(maxsize=100)
+def cached_generate_response(prompt, tokenizer, model):
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
+    outputs = model.generate(inputs["input_ids"], max_new_tokens=50, temperature=0.7, do_sample=True)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
-# Function to generate response with debugging logs
+# Function to generate response with optimization
 def generate_response(prompt):
     try:
-        # Tokenize input prompt with dynamic padding and truncation
-        log_box.text_area("Debugging Logs", "Tokenizing the prompt...", height=200)
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
-
-        # Display tokenized inputs
-        log_box.text_area("Debugging Logs", f"Tokenized inputs: {inputs['input_ids']}", height=200)
-        
-        # Generate output from model
-        log_box.text_area("Debugging Logs", "Generating output...", height=200)
-        outputs = model.generate(inputs["input_ids"], max_new_tokens=100, temperature=0.7, do_sample=True)
-
-        # Display the raw output from the model
-        log_box.text_area("Debugging Logs", f"Raw model output (tokens): {outputs}", height=200)
-
-        # Decode and return response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Display the final decoded response
-        log_box.text_area("Debugging Logs", f"Decoded response: {response}", height=200)
-        
-        return response
+        # Check cache for previous result (for repeated queries)
+        cached_response = cached_generate_response(prompt, tokenizer, model)
+        return cached_response
     except Exception as e:
         st.error(f"Error during text generation: {e}")
         return "Sorry, I couldn't process your request."
@@ -93,7 +76,6 @@ if user_input:
             response = generate_response(user_input)
             st.write(response)
 
-    # Update session state
+    # Update session state with new messages
     st.session_state.messages.append({"role": "user", "content": user_input})
     st.session_state.messages.append({"role": "assistant", "content": response})
-