Update app.py

app.py

@@ -6,7 +6,6 @@ from datetime import datetime, timedelta, timezone
 import torch
 from config import hugging_face_token, init_google_sheets_client, models, default_model_name, user_names, google_sheets_name, MAX_INTERACTIONS
 import spaces
-import gspread
 
 # Hack for ZeroGPU
 torch.jit.script = lambda f: f
@@ -66,7 +65,7 @@ def load_model(model_name):
 # Ensure the initial model is loaded
 tokenizer, model = load_model(selected_model)
 
-# Chat history and interaction counter
+# Chat history and interaction count
 chat_history = []
 interaction_count = 0
 
@@ -78,12 +77,6 @@ def interact(user_input, history, interaction_count):
         if tokenizer is None or model is None:
             raise ValueError("Tokenizer or model is not initialized.")
         
-        interaction_count += 1
-        print(f"Interaction count: {interaction_count}")
-
-        if interaction_count >= MAX_INTERACTIONS:
-            user_input += ". Thank you for the questions. That's all for now. Goodbye!"
-
         messages = history + [{"role": "user", "content": user_input}]
         
         # Ensure roles alternate correctly
@@ -95,11 +88,16 @@ 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)  # 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})
+        
+        # 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})
         
         formatted_history = [(entry["content"], None) if entry["role"] == "user" else (None, entry["content"]) for entry in history if entry["role"] in ["user", "assistant"]]
@@ -112,8 +110,11 @@ 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, interaction_count
+    global selected_story
+    global data  # Ensure data is reset
     data = []  # Reset data for new story
+    interaction_count = 0  # Reset interaction count
     tokenizer, model = load_model(model_name)
     selected_story = title
     for story in stories:
@@ -131,9 +132,10 @@ Here is the story:
                 chat_history.append({"role": "system", "content": combined_message})
 
                 # Generate the first question based on the story
-                _, formatted_history, chat_history, interaction_count = interact("Please ask a simple question about the story to encourage interaction.", chat_history, 0)
+                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"]  # Reset the data table and return the story
             else:
                 print("Combined message is empty.")
         else:
@@ -192,7 +194,7 @@ with gr.Blocks() as demo:
     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")
+    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)
@@ -214,10 +216,9 @@ 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_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_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])
+    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)])
     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()