Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import pandas as pd
 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
+from config import hugging_face_token, init_google_sheets_client, models, default_model_name, user_names, google_sheets_name, max_interactions
 import spaces
 
 # Hack for ZeroGPU
@@ -13,8 +13,8 @@ torch.jit.script = lambda f: f
 # Initialize Google Sheets client
 client = init_google_sheets_client()
 sheet = client.open(google_sheets_name)
-stories_sheet = sheet.worksheet("Stories")  # Assuming stories are in a separate sheet
-prompts_sheet = sheet.worksheet("System Prompts")  # Assuming system prompts are in a separate sheet
+stories_sheet = sheet.worksheet("Stories")  # Assuming stories are in the second sheet (index 1)
+system_prompts_sheet = sheet.worksheet("System Prompts")  # Assuming system prompts are in a separate sheet
 
 # Load stories from Google Sheets
 def load_stories():
@@ -23,14 +23,14 @@ def load_stories():
     return stories
 
 # Load system prompts from Google Sheets
-def load_prompts():
-    prompts_data = prompts_sheet.get_all_values()
-    prompts = [prompt[0] for prompt in prompts_data if prompt[0] != "System Prompts"]  # Skip header row
-    return prompts
+def load_system_prompts():
+    system_prompts_data = system_prompts_sheet.get_all_values()
+    system_prompts = [prompt[0] for prompt in system_prompts_data[1:]]  # Skip header row
+    return system_prompts
 
-# Load available stories and prompts
+# Load available stories and system prompts
 stories = load_stories()
-prompts = load_prompts()
+system_prompts = load_system_prompts()
 
 # Initialize the selected model
 selected_model = default_model_name
@@ -65,21 +65,18 @@ def load_model(model_name):
 # Ensure the initial model is loaded
 tokenizer, model = load_model(selected_model)
 
-# Chat history and interaction count
+# Chat history and interaction counter
 chat_history = []
+interaction_count = 0
 
 # Function to handle interaction with model
 @spaces.GPU
-def interact(user_input, history, interaction_count):
-    global tokenizer, model
+def interact(user_input, history):
+    global tokenizer, model, interaction_count
     try:
         if tokenizer is None or model is None:
             raise ValueError("Tokenizer or model is not initialized.")
         
-        # Concatenate a final message if max interactions are reached
-        if interaction_count >= MAX_INTERACTIONS - 1:
-            user_input += ". Thank you for the questions. That's all for now. Goodbye!"
-
         messages = history + [{"role": "user", "content": user_input}]
         
         # Ensure roles alternate correctly
@@ -89,10 +86,6 @@ 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
-        
         # 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
@@ -101,9 +94,16 @@ def interact(user_input, history, interaction_count):
         # 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}")
+        
+        if interaction_count >= max_interactions:
+            response += ". Thank you for the questions. That's all for now. Goodbye!"
+            history[-1]["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"]]
-        return "", formatted_history, history, interaction_count
+        return "", formatted_history, history
     except Exception as e:
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -112,16 +112,13 @@ 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
-    global selected_story
-    global data  # Ensure data is reset
+    global chat_history, selected_story, data, interaction_count
     data = []  # Reset data for new story
+    interaction_count = 0  # Reset interaction counter
     tokenizer, model = load_model(model_name)
     selected_story = title
-    story_text = ""
     for story in stories:
         if story["title"] == title:
-            story_text = story["story"]
             system_prompt = f"""
 {system_prompt}
 Here is the story:
@@ -136,9 +133,9 @@ Here is the story:
 
                 # Generate the first question based on 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, 0)
+                _, formatted_history, chat_history = interact(question_prompt, chat_history)
 
-                return formatted_history, chat_history, gr.update(value=[]), gr.update(value=story_text), interaction_count  # Reset the data table and update the selected story textbox
+                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:
@@ -178,11 +175,15 @@ def save_comment_score(chat_responses, score, comment, story_name, user_name, sy
     ])
 
     # Append data to Google Sheets
-    sheet = client.open(google_sheets_name).sheet1  # Assuming results are saved in sheet1
-    sheet.append_row([timestamp_str, user_name, model_name, system_prompt, story_name, last_user_message, last_assistant_message, score, comment])
+    try:
+        user_sheet = client.open(google_sheets_name).worksheet(user_name)
+    except gspread.exceptions.WorksheetNotFound:
+        user_sheet = client.open(google_sheets_name).add_worksheet(title=user_name, rows="100", cols="20")
+    
+    user_sheet.append_row([timestamp_str, user_name, model_name, system_prompt, story_name, last_user_message, last_assistant_message, score, comment])
 
     df = pd.DataFrame(data, columns=["Timestamp", "User Name", "Model Name", "System Prompt", "Story Name", "User Input", "Chat Response", "Score", "Comment"])
-    return df, gr.update(value="")  # Clear the comment input box
+    return df[["User Input", "Chat Response", "Score", "Comment"]], gr.update(value="")  # Show only the required columns and clear the comment input box
 
 # Create the chat interface using Gradio Blocks
 with gr.Blocks() as demo:
@@ -192,11 +193,11 @@ with gr.Blocks() as demo:
     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=prompts, label="Select System Prompt")
 
-    send_story_button = gr.Button("Send Story")
+    system_prompt_dropdown = gr.Dropdown(choices=system_prompts, label="Select System Prompt")
 
-    selected_story_textbox = gr.Textbox(label="Selected Story", interactive=False)
+    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):
@@ -215,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], outputs=[chatbot_input, chatbot_output, chat_history_json])
     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()