Update app.py

app.py

@@ -150,68 +150,17 @@ def push_feedback_to_hub(hf_token=None):
         print(f"Error pushing feedback data to Hub: {e}")
         return False
 
-# Modified predict function to update conversation state
-@spaces.GPU(duration=120)
-def predict(message, history, state, temperature, top_p):
-    # Create a deep copy of history to ensure we don't modify the original
-    current_history = history.copy()
-    
-    # Update history with user message
-    current_history.append({"role": "user", "content": message})
-    
-    # Update the conversation state with user message
-    if not state:
-        state = []
-    state = current_history.copy()
-    
-    input_text = tokenizer.apply_chat_template(current_history, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
-    
-    # Create a streamer
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    
-    # Set up generation parameters
-    generation_kwargs = {
-        "input_ids": inputs,
-        "max_new_tokens": 1024,
-        "temperature": float(temperature),
-        "top_p": float(top_p),
-        "do_sample": True,
-        "streamer": streamer,
-        "eos_token_id": 128009,
-    }
-    
-    # Run generation in a separate thread
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    
-    # Yield from the streamer as tokens are generated
-    partial_text = ""
-    for new_text in streamer:
-        partial_text += new_text
-        
-        # Create a temporary state with partial response
-        temp_history = current_history.copy()
-        temp_history.append({"role": "assistant", "content": partial_text})
-        temp_state = temp_history.copy()
-        
-        yield partial_text, temp_state
-    
-    # After full generation, update state with assistant's final response
-    current_history.append({"role": "assistant", "content": partial_text})
-    state = current_history.copy()
-    
-    # Print debug info
-    print(f"Updated state with {len(state)} messages")
-    print(f"Last message: {state[-1]['role']}: {state[-1]['content'][:30]}...")
-    
-    return partial_text, state
-
 # Function to handle the research feedback submission
-def submit_research_feedback(conversation_state, satisfaction, feedback_text):
+def submit_research_feedback(conv_history, satisfaction, feedback_text):
     """Save user feedback both locally and to HuggingFace Hub"""
+    # Print debug information
+    print(f"Saving feedback with conversation history containing {len(conv_history)} messages")
+    if conv_history and len(conv_history) > 0:
+        print(f"First message: {conv_history[0]['role']}: {conv_history[0]['content'][:30]}...")
+        print(f"Last message: {conv_history[-1]['role']}: {conv_history[-1]['content'][:30]}...")
+    
     # Save locally first
-    feedback_id = save_feedback_locally(conversation_state, satisfaction, feedback_text)
+    feedback_id = save_feedback_locally(conv_history, satisfaction, feedback_text)
     
     # Get token from environment variable
     env_token = os.environ.get("HF_TOKEN")
@@ -226,25 +175,111 @@ def submit_research_feedback(conversation_state, satisfaction, feedback_text):
     
     return status_msg
 
+# Initial state - set up at app start
+def initialize_state():
+    """Initialize the conversation state - this could load previous sessions or start fresh"""
+    return []  # Start with empty conversation history
+
 # Create the Gradio blocks interface
 with gr.Blocks() as demo:
-    # State to track conversation history
-    conversation_state = gr.State([])
+    # Create state to store full conversation history with proper initialization
+    conv_state = gr.State(initialize_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
-                    conversation_state.value = updated_state
-                    yield partial_response
+            # Create a custom predict function that updates our state
+            def enhanced_predict(message, history, temperature, top_p, state):
+                # Initialize state if needed
+                if state is None:
+                    state = []
+                    print("Initializing empty state")
+                
+                # Copy history to state if state is empty but history exists
+                if len(state) == 0 and len(history) > 0:
+                    state = history.copy()
+                    print(f"Copied {len(history)} messages from history to state")
+                
+                # Add user message to state
+                state.append({"role": "user", "content": message})
+                
+                # Process with the model (this doesn't modify the original history)
+                input_text = tokenizer.apply_chat_template(state, tokenize=False, add_generation_prompt=True)
+                inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+                
+                # Create a streamer
+                streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+                
+                # Set up generation parameters
+                generation_kwargs = {
+                    "input_ids": inputs,
+                    "max_new_tokens": 1024,
+                    "temperature": float(temperature),
+                    "top_p": float(top_p),
+                    "do_sample": True,
+                    "streamer": streamer,
+                    "eos_token_id": 128009,
+                }
+                
+                # Run generation in a separate thread
+                thread = Thread(target=model.generate, kwargs=generation_kwargs)
+                thread.start()
+                
+                # Yield from the streamer as tokens are generated
+                response = ""
+                for new_text in streamer:
+                    response += new_text
+                    # For each partial response, yield the text only
+                    # We'll update the state after generation is complete
+                    yield response
+                
+                # After generation completes, update our state with the final response
+                state.append({"role": "assistant", "content": response})
+                
+                # Return the updated state
+                return state
+            
+            # Create a wrapper that connects to ChatInterface but also updates our state
+            def chat_with_state(message, history, temperature, top_p):
+                # This function is what interfaces with the ChatInterface
+                nonlocal conv_state
+                
+                # Access the current state
+                current_state = conv_state.value if conv_state.value else []
+                
+                # Call the main function that generates responses and updates state
+                # This is a generator function, so we need to iterate through its outputs
+                response_gen = enhanced_predict(message, history, temperature, top_p, current_state)
+                
+                # For each response, yield it and also update our state at the end
+                last_response = None
+                for response in response_gen:
+                    last_response = response
+                    yield response
+                
+                # After generation is complete, update our state
+                if last_response is not None:
+                    # Create a full copy of the history plus the new exchange
+                    updated_state = []
+                    # Add all previous history
+                    for msg in history:
+                        updated_state.append(msg.copy())
+                    # Add new exchange
+                    updated_state.append({"role": "user", "content": message})
+                    updated_state.append({"role": "assistant", "content": last_response})
+                    
+                    # Store in our state
+                    conv_state.value = updated_state
+                    
+                    # Debug
+                    print(f"Updated conversation state with {len(updated_state)} messages")
+                    if updated_state:
+                        last_msg = updated_state[-1]
+                        print(f"Last message: {last_msg['role']}: {last_msg['content'][:30]}...")
             
             # Create ChatInterface
             chatbot = gr.ChatInterface(
                 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")
                 ],
@@ -282,10 +317,10 @@ with gr.Blocks() as demo:
         feedback_modal
     )
     
-    # Connect the submit button to the submit_research_feedback function with the current conversation state
+    # Connect the submit button to the submit_research_feedback function
     submit_button.click(
         submit_research_feedback,
-        inputs=[conversation_state, satisfaction, feedback_text],
+        inputs=[conv_state, satisfaction, feedback_text],
         outputs=response_text
     )