Update app.py

app.py

@@ -67,7 +67,6 @@ tokenizer, model = load_model(selected_model)
 
 # Chat history
 chat_history = []
-interaction_count = 0  # Initialize interaction count
 
 # Function to handle interaction with model
 @spaces.GPU
@@ -77,11 +76,9 @@ 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
-
+        interaction_count += 1
+        print(f"Interaction count: {interaction_count}")
+        
         messages = history + [{"role": "user", "content": user_input}]
         
         # Ensure roles alternate correctly
@@ -89,26 +86,22 @@ def interact(user_input, history, interaction_count):
             if messages[i-1].get("role") == messages[i].get("role"):
                 raise ValueError("Conversation roles must alternate user/assistant/user/assistant/...")
         
+        if interaction_count >= MAX_INTERACTIONS:
+            user_input += ". Thank you for the questions. That's all for now. Goodbye!"
+        
         prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         
         # 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
+        chat_history_ids = model.generate(input_ids, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
         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})
         
-        interaction_count += 1
-        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():
@@ -121,9 +114,8 @@ def send_selected_story(title, model_name, system_prompt):
     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
+    interaction_count = 0  # Reset interaction count for new story
     tokenizer, model = load_model(model_name)
     selected_story = title
     for story in stories:
@@ -144,7 +136,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"], interaction_count  # Reset the data table and return the story
+                return formatted_history, chat_history, gr.update(value=[]), story["story"]
             else:
                 print("Combined message is empty.")
         else:
@@ -196,38 +188,39 @@ 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")
-
-    gr.Markdown("## 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")
-
-    gr.Markdown("## 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")
-    
-    gr.Markdown("## 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"])
+    with gr.Tabs():
+        with gr.TabItem("Chat"):
+            gr.Markdown("# Chat with Model")
+
+            with gr.Group(label="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.Group(label="Chat"):
+                selected_story_textbox = gr.Textbox(label="Selected Story", lines=10, interactive=False)
+                chatbot_input = gr.Textbox(placeholder="Type your message here...", label="User Input")
+                send_message_button = gr.Button("Send")
+                chatbot_output = gr.Chatbot(label="Chat History")
+
+            with gr.Group(label="Evaluation"):
+                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.TabItem("Instructions"):
+            gr.Markdown("# Instructions")
+            instructions_textbox = gr.Textbox(value="Here are the instructions on how to use the application...", lines=20)
 
     chat_history_json = gr.JSON(value=[], visible=False)
-    interaction_counter_json = gr.JSON(value=0, 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, interaction_counter_json])
-    send_message_button.click(fn=interact, inputs=[chatbot_input, chat_history_json, interaction_counter_json], outputs=[chatbot_input, chatbot_output, chat_history_json, interaction_counter_json])
+    interaction_count = gr.Number(value=0, 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, interaction_count], outputs=[chatbot_input, chatbot_output, chat_history_json, interaction_count])
     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()