Update app.py

app.py

@@ -17,21 +17,23 @@ def load_model():
         st.error(f"Failed to load model: {e}")
         return None, None
 
+# Load model and tokenizer
 tokenizer, model = load_model()
 
 if not tokenizer or not model:
     st.stop()
 
-# Default to CPU
-device = torch.device("cpu")
-if model is not None:
-    model = model.to(device)
+# Default to CPU, or use GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
 
-if "messages" not in st.session_state.keys():
+# Initialize session state messages if not already initialized
+if "messages" not in st.session_state:
     st.session_state.messages = [
         {"role": "assistant", "content": "Hello there! How can I help you with gardening today?"}
     ]
 
+# Display the conversation history
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.write(message["content"])
@@ -39,11 +41,11 @@ for message in st.session_state.messages:
 def generate_response(prompt):
     try:
         # Tokenize the input prompt
-        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
+
+        # Ensure the model is generating properly (without a target)
+        outputs = model.generate(inputs["input_ids"], max_new_tokens=150, temperature=0.7, do_sample=True)
 
-        # Make sure that inputs are passed properly to the model
-        outputs = model.generate(inputs["input_ids"], max_new_tokens=150, temperature=0.7)
-        
         # Decode the output to text
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response
@@ -51,7 +53,7 @@ def generate_response(prompt):
         st.error(f"Error during text generation: {e}")
         return "Sorry, I couldn't process your request."
 
-
+# User input field for asking questions
 user_input = st.chat_input("Type your gardening question here:")
 
 if user_input:
@@ -59,10 +61,12 @@ if user_input:
     with st.chat_message("user"):
         st.write(user_input)
 
-with st.chat_message("assistant"):
+    # Generate and display assistant's response
+    with st.chat_message("assistant"):
         with st.spinner("I'm gonna tell you..."):
             response = generate_response(user_input)
             st.write(response)
 
-st.session_state.messages.append({"role": "user", "content": user_input})
-st.session_state.messages.append({"role": "assistant", "content": response})
+    # Update session state with the new conversation
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    st.session_state.messages.append({"role": "assistant", "content": response})