Update app.py

app.py

@@ -65,9 +65,9 @@ def load_model(model_name):
 # Ensure the initial model is loaded
 tokenizer, model = load_model(selected_model)
 
-# Chat history and interaction count
+# Chat history
 chat_history = []
-interaction_count = 0
+interaction_count = 0  # Initialize interaction count
 
 # Function to handle interaction with model
 @spaces.GPU
@@ -77,6 +77,11 @@ def interact(user_input, history, interaction_count):
         if tokenizer is None or model is None:
             raise ValueError("Tokenizer or model is not initialized.")
         
+        if interaction_count >= MAX_INTERACTIONS:
+            history.append({"role": "assistant", "content": "Thank you for the conversation! Have a great day!"})
+            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
+
         messages = history + [{"role": "user", "content": user_input}]
         
         # Ensure roles alternate correctly
@@ -88,31 +93,35 @@ def interact(user_input, 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, temperature=0.1)
+        chat_history_ids = model.generate(input_ids, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)  # Increase max_new_tokens
         response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
         
         # Update chat history with generated response
         history.append({"role": "user", "content": user_input})
+        history.append({"role": "assistant", "content": response})
         
-        # Check if it's the last interaction
         interaction_count += 1
-        if interaction_count >= MAX_INTERACTIONS:
-            response += ". Thank you for the questions. That's all for now. Goodbye!"
-        history.append({"role": "assistant", "content": response})
+        print(f"Interaction count: {interaction_count}")
         
         formatted_history = [(entry["content"], None) if entry["role"] == "user" else (None, entry["content"]) for entry in history if entry["role"] in ["user", "assistant"]]
+        
+        if interaction_count >= MAX_INTERACTIONS:
+            history.append({"role": "assistant", "content": "Thank you for the conversation! Have a great day!"})
+            formatted_history.append((None, "Thank you for the conversation! Have a great day!"))
+        
         return "", formatted_history, history, interaction_count
     except Exception as e:
-        if torch.cuda.is_available():
+        if torch.cuda.is available():
             torch.cuda.empty_cache()
         print(f"Error during interaction: {e}")
         raise gr.Error(f"An error occurred during interaction: {str(e)}")
 
 # Function to send selected story and initial message
 def send_selected_story(title, model_name, system_prompt):
-    global chat_history, interaction_count
+    global chat_history
     global selected_story
     global data  # Ensure data is reset
+    global interaction_count
     data = []  # Reset data for new story
     interaction_count = 0  # Reset interaction count
     tokenizer, model = load_model(model_name)
@@ -135,7 +144,7 @@ Here is the story:
                 question_prompt = "Please ask a simple question about the story to encourage interaction."
                 _, formatted_history, chat_history, interaction_count = interact(question_prompt, chat_history, interaction_count)
 
-                return formatted_history, chat_history, gr.update(value=[]), story["story"]  # Reset the data table and return the story
+                return formatted_history, chat_history, gr.update(value=[]), story["story"], interaction_count  # Reset the data table and return the story
             else:
                 print("Combined message is empty.")
         else:
@@ -188,37 +197,36 @@ def save_comment_score(chat_responses, score, comment, story_name, user_name, sy
 # Create the chat interface using Gradio Blocks
 with gr.Blocks() as demo:
     gr.Markdown("# Chat with Model")
-
-    model_dropdown = gr.Dropdown(choices=list(models.keys()), label="Select Model", value=selected_model)
-    user_dropdown = gr.Dropdown(choices=user_names, label="Select User Name")
-    initial_story = stories[0]["title"] if stories else None
-    story_dropdown = gr.Dropdown(choices=[story["title"] for story in stories], label="Select Story", value=initial_story)
-
-    system_prompt_dropdown = gr.Dropdown(choices=system_prompts, label="Select System Prompt", value=system_prompts[0])
-
-    send_story_button = gr.Button("Send Story")
-    selected_story_textbox = gr.Textbox(label="Selected Story", lines=10, interactive=False)
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            chatbot_input = gr.Textbox(placeholder="Type your message here...", label="User Input")
-            send_message_button = gr.Button("Send")
-
-        with gr.Column(scale=2):
-            chatbot_output = gr.Chatbot(label="Chat History")
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            score_input = gr.Slider(minimum=0, maximum=5, step=1, label="Score")
-            comment_input = gr.Textbox(placeholder="Add a comment...", label="Comment")
-            save_button = gr.Button("Save Score and Comment")
-
-    data_table = gr.DataFrame(headers=["User Input", "Chat Response", "Score", "Comment"])
+    
+    with gr.Tab("Context"):
+        model_dropdown = gr.Dropdown(choices=list(models.keys()), label="Select Model", value=selected_model)
+        user_dropdown = gr.Dropdown(choices=user_names, label="Select User Name")
+        initial_story = stories[0]["title"] if stories else None
+        story_dropdown = gr.Dropdown(choices=[story["title"] for story in stories], label="Select Story", value=initial_story)
+        system_prompt_dropdown = gr.Dropdown(choices=system_prompts, label="Select System Prompt", value=system_prompts[0])
+        send_story_button = gr.Button("Send Story")
+    
+    with gr.Tab("Chat"):
+        selected_story_textbox = gr.Textbox(label="Selected Story", lines=10, interactive=False)
+        with gr.Row():
+            with gr.Column(scale=1):
+                chatbot_input = gr.Textbox(placeholder="Type your message here...", label="User Input")
+                send_message_button = gr.Button("Send")
+            with gr.Column(scale=2):
+                chatbot_output = gr.Chatbot(label="Chat History")
+    
+    with gr.Tab("Evaluation"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                score_input = gr.Slider(minimum=0, maximum=5, step=1, label="Score")
+                comment_input = gr.Textbox(placeholder="Add a comment...", label="Comment")
+                save_button = gr.Button("Save Score and Comment")
+        data_table = gr.DataFrame(headers=["User Input", "Chat Response", "Score", "Comment"])
 
     chat_history_json = gr.JSON(value=[], visible=False)
-
-    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])
-    send_message_button.click(fn=interact, inputs=[chatbot_input, chat_history_json, gr.State(interaction_count)], outputs=[chatbot_input, chatbot_output, chat_history_json, gr.State(interaction_count)])
+    
+    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, gr.JSON(value=0)])
+    send_message_button.click(fn=interact, inputs=[chatbot_input, chat_history_json, gr.JSON(value=0)], outputs=[chatbot_input, chatbot_output, chat_history_json, gr.JSON(value=0)])
     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()