Update app.py

app.py

@@ -3,7 +3,7 @@ import streamlit as st
 from models import load_model
 
 # Load the model once
-demo = load_model()
+generator = load_model()
 
 # Page configuration
 st.set_page_config(
@@ -73,20 +73,25 @@ if prompt := st.chat_input("Type your message..."):
     try:
         # Generate response using the model
         with st.spinner("Generating response..."):
-            # Pass inputs as positional arguments to the Gradio model
-            response = demo(
-                f"{system_message}\n\nUser: {prompt}\nAssistant:",
-                max_length=max_tokens,  # Gradio parameter
+            full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
+            response = generator(
+                full_prompt,
+                max_length=max_tokens,
                 temperature=temperature,
-                top_p=top_p
-            )
+                top_p=top_p,
+                do_sample=True,
+                num_return_sequences=1
+            )[0]['generated_text']
+            
+            # Extract only the assistant's response
+            assistant_response = response.split("Assistant:")[-1].strip()
         
         # Display assistant response
         with st.chat_message("assistant"):
-            st.markdown(response)
+            st.markdown(assistant_response)
         
         # Add assistant response to chat history
-        st.session_state.messages.append({"role": "assistant", "content": response})
+        st.session_state.messages.append({"role": "assistant", "content": assistant_response})
     
     except Exception as e:
-        st.error(f"An error occurred: {str(e)}")
+        st.error(f"An error occurred: {str(e)}")