Update app.py

app.py

@@ -93,16 +93,14 @@ def push_feedback_to_hub(hf_token=None):
         print(f"Error pushing feedback data to Hub: {e}")
         return False
 
-# Create a State to store chat history
-chat_history_state = []
-
+# Modified predict function to update conversation state
 @spaces.GPU(duration=120)
-def predict(message, history, temperature, top_p):
-    global chat_history_state
-    
-    # Update our chat history state
+def predict(message, history, state, temperature, top_p):
+    # Update history with user message
     history.append({"role": "user", "content": message})
-    chat_history_state = history.copy()
+    
+    # Update the conversation state
+    state = history.copy()
     
     input_text = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
@@ -129,18 +127,18 @@ def predict(message, history, temperature, top_p):
     partial_text = ""
     for new_text in streamer:
         partial_text += new_text
-        yield partial_text
+        yield partial_text, state
     
-    # After generation is complete, update chat history state with the assistant response
-    chat_history_state.append({"role": "assistant", "content": partial_text})
+    # After full generation, update state with assistant's response
+    history.append({"role": "assistant", "content": partial_text})
+    state = history.copy()
+    return partial_text, state
 
 # Function to handle the research feedback submission
-def submit_research_feedback(satisfaction, feedback_text):
+def submit_research_feedback(conversation_state, satisfaction, feedback_text):
     """Save user feedback both locally and to HuggingFace Hub"""
-    global chat_history_state
-    
     # Save locally first
-    feedback_id = save_feedback_locally(chat_history_state, satisfaction, feedback_text)
+    feedback_id = save_feedback_locally(conversation_state, satisfaction, feedback_text)
     
     # Get token from environment variable
     env_token = os.environ.get("HF_TOKEN")
@@ -155,16 +153,29 @@ def submit_research_feedback(satisfaction, feedback_text):
     
     return status_msg
 
-# Create the Gradio interface
+# Create the Gradio blocks interface
 with gr.Blocks() as demo:
+    # State to track conversation history
+    conversation_state = gr.State([])
+    
     with gr.Row():
         with gr.Column(scale=3):
+            # Custom chat function wrapper to update state
+            def chat_with_state(message, history, state, temperature, top_p):
+                for partial_response, updated_state in predict(message, history, state, temperature, top_p):
+                    # Update our state with each yield
+                    state = updated_state
+                    yield partial_response, state
+            
+            # Create ChatInterface
             chatbot = gr.ChatInterface(
-                predict,
+                chat_with_state,
                 additional_inputs=[
+                    conversation_state,
                     gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
                     gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
                 ],
+                additional_outputs=[conversation_state],
                 type="messages"
             )
         
@@ -199,10 +210,10 @@ with gr.Blocks() as demo:
         feedback_modal
     )
     
-    # Connect the submit button to the submit_research_feedback function with the current chat history
+    # Connect the submit button to the submit_research_feedback function with the current conversation state
     submit_button.click(
-        lambda satisfaction, feedback_text: submit_research_feedback(satisfaction, feedback_text),
-        inputs=[satisfaction, feedback_text],
+        submit_research_feedback,
+        inputs=[conversation_state, satisfaction, feedback_text],
         outputs=response_text
     )