Update app.py

app.py

@@ -85,6 +85,15 @@ def interact(user_input, history, interaction_count):
         
         prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         
+        # Check if the maximum number of interactions has been reached
+        interaction_count += 1
+        print(f"Interaction count: {interaction_count}")  # Print the interaction count
+        if interaction_count >= MAX_INTERACTIONS:
+            farewell_message = "Thank you for the conversation! Have a great day!"
+            history.append({"role": "assistant", "content": farewell_message})
+            formatted_history = [(entry["content"], None) if entry["role"] == "user" else (None, entry["content"]) for entry in history if entry["role"] in ["user", "assistant"]]
+            return "", formatted_history, history, interaction_count
+        
         # Generate response using selected model
         input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to("cuda")
         chat_history_ids = model.generate(input_ids, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)  # Increase max_new_tokens
@@ -94,13 +103,6 @@ def interact(user_input, history, interaction_count):
         history.append({"role": "user", "content": user_input})
         history.append({"role": "assistant", "content": response})
 
-        # Check if the maximum number of interactions has been reached
-        interaction_count += 1
-        print(f"Interaction count: {interaction_count}")  # Print the interaction count
-        if interaction_count >= MAX_INTERACTIONS:
-            farewell_message = "Thank you for the conversation! Have a great day!"
-            history.append({"role": "assistant", "content": farewell_message})
-        
         formatted_history = [(entry["content"], None) if entry["role"] == "user" else (None, entry["content"]) for entry in history if entry["role"] in ["user", "assistant"]]
         return "", formatted_history, history, interaction_count
     except Exception as e:
@@ -111,7 +113,9 @@ def interact(user_input, history, interaction_count):
 
 # Function to send selected story and initial message
 def send_selected_story(title, model_name, system_prompt):
-    global chat_history, selected_story, data
+    global chat_history
+    global selected_story
+    global data  # Ensure data is reset
     data = []  # Reset data for new story
     tokenizer, model = load_model(model_name)
     selected_story = title
@@ -212,10 +216,10 @@ with gr.Blocks() as demo:
     data_table = gr.DataFrame(headers=["User Input", "Chat Response", "Score", "Comment"])
 
     chat_history_json = gr.JSON(value=[], visible=False)
-    interaction_state = gr.State(0)
+    interaction_count_state = gr.State(0)
 
-    send_story_button.click(fn=send_selected_story, inputs=[story_dropdown, model_dropdown, system_prompt_dropdown], outputs=[chatbot_output, chat_history_json, data_table, selected_story_textbox, interaction_state])
-    send_message_button.click(fn=interact, inputs=[chatbot_input, chat_history_json, interaction_state], outputs=[chatbot_input, chatbot_output, chat_history_json, interaction_state])
+    send_story_button.click(fn=send_selected_story, inputs=[story_dropdown, model_dropdown, system_prompt_dropdown], outputs=[chatbot_output, chat_history_json, data_table, selected_story_textbox, interaction_count_state])
+    send_message_button.click(fn=interact, inputs=[chatbot_input, chat_history_json, interaction_count_state], outputs=[chatbot_input, chatbot_output, chat_history_json, interaction_count_state])
     save_button.click(fn=save_comment_score, inputs=[chatbot_output, score_input, comment_input, story_dropdown, user_dropdown, system_prompt_dropdown], outputs=[data_table, comment_input])
 
 demo.launch()