Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 import os
 from dotenv import load_dotenv
 
@@ -11,22 +12,33 @@ 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 the pipeline
-def load_pipeline():
+# Function to load model only once
+def load_model():
     try:
-        # Create the text-generation pipeline using the model from Hugging Face
-        pipe = pipeline("text-generation", model="TheSheBots/UrbanGardening", use_auth_token=HF_TOKEN)
-        return pipe
+        # 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("TheSheBots/UrbanGardening", 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
     except Exception as e:
-        st.error(f"Failed to load model pipeline: {e}")
-        return None
+        st.error(f"Failed to load model: {e}")
+        return None, None
 
-# Load the pipeline (cached)
-pipe = load_pipeline()
+# Load model and tokenizer (cached)
+tokenizer, model = load_model()
 
-if not pipe:
+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")
+model = model.to(device)
+
 # Initialize session state messages
 if "messages" not in st.session_state:
     st.session_state.messages = [
@@ -38,6 +50,37 @@ 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()
+
+# Function to generate response with debugging logs
+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
+    except Exception as e:
+        st.error(f"Error during text generation: {e}")
+        return "Sorry, I couldn't process your request."
+
 # User input field for gardening questions
 user_input = st.chat_input("Type your gardening question here:")
 
@@ -47,15 +90,10 @@ if user_input:
 
     with st.chat_message("assistant"):
         with st.spinner("Generating your answer..."):
-            try:
-                # Generate the response using the pipeline
-                response = pipe(user_input, max_length=150, num_return_sequences=1, temperature=0.7)[0]['generated_text']
-                st.write(response)
-            except Exception as e:
-                st.error(f"Error during text generation: {e}")
-                response = "Sorry, I couldn't process your request."
-                st.write(response)
+            response = generate_response(user_input)
+            st.write(response)
 
     # Update session state
     st.session_state.messages.append({"role": "user", "content": user_input})
     st.session_state.messages.append({"role": "assistant", "content": response})
+