Update app.py

app.py

@@ -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")
-prompts_sheet = sheet.worksheet("System Prompts")
+stories_sheet = sheet.get_worksheet(1)  # Assuming stories are in the second sheet (index 1)
+prompts_sheet = sheet.get_worksheet(2)  # Assuming system prompts are in the third sheet (index 2)
 
 # Load stories from Google Sheets
 def load_stories():
@@ -22,15 +22,17 @@ def load_stories():
     stories = [{"title": story[0], "story": story[1]} for story in stories_data if story[0] != "Title"]  # Skip header row
     return stories
 
+# Load available stories
+stories = load_stories()
+
 # Load system prompts from Google Sheets
-def load_prompts():
+def load_system_prompts():
     prompts_data = prompts_sheet.get_all_values()
-    prompts = [prompt[0] for prompt in prompts_data[1:] if prompt[0] != "System Prompt"]  # Skip header row
-    return prompts
+    system_prompts = [row[0] for row in prompts_data if row[0] != "System Prompt"]  # Skip header row
+    return system_prompts
 
-# Load available stories and prompts
-stories = load_stories()
-prompts = load_prompts()
+# Load available system prompts
+system_prompts = load_system_prompts()
 
 # Initialize the selected model
 selected_model = default_model_name
@@ -106,12 +108,10 @@ def interact(user_input, history):
 def send_selected_story(title, model_name, system_prompt):
     global chat_history
     global selected_story
-    global selected_system_prompt
     global data  # Ensure data is reset
     data = []  # Reset data for new story
     tokenizer, model = load_model(model_name)
     selected_story = title
-    selected_system_prompt = system_prompt
     for story in stories:
         if story["title"] == title:
             system_prompt = f"""
@@ -130,14 +130,14 @@ Here is the story:
                 question_prompt = "Please ask a simple question about the story to encourage interaction."
                 _, formatted_history, chat_history = interact(question_prompt, chat_history)
 
-                return formatted_history, chat_history, gr.update(value=[], visible=True), gr.update(value=story['story'])  # Reset the data table and show selected story
+                return formatted_history, chat_history, gr.update(value=[]), gr.update(value=story['story'])  # Reset the data table and update selected story
             else:
                 print("Combined message is empty.")
         else:
             print("Story title does not match.")
 
 # Function to save comment and score
-def save_comment_score(chat_responses, score, comment, story_name, user_name):
+def save_comment_score(chat_responses, score, comment, story_name, user_name, system_prompt):
     last_user_message = ""
     last_assistant_message = ""
 
@@ -155,7 +155,6 @@ def save_comment_score(chat_responses, score, comment, story_name, user_name):
     timestamp = datetime.now(timezone.utc) - timedelta(hours=3)  # Adjust to GMT-3
     timestamp_str = timestamp.strftime("%Y-%m-%d %H:%M:%S")
     model_name = selected_model
-    system_prompt = selected_system_prompt
 
     # Append data to local data storage
     data.append([
@@ -171,11 +170,15 @@ def save_comment_score(chat_responses, score, comment, story_name, user_name):
     ])
 
     # Append data to Google Sheets
-    sheet = client.open(google_sheets_name).worksheet(user_name)  # Save to user-specific sheet
+    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])
 
+    # Create a DataFrame for display
+    display_data = [[last_user_message, last_assistant_message, score, comment]]
+    df_display = pd.DataFrame(display_data, columns=["User Input", "Chat Response", "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_display, gr.update(value="")  # Clear the comment input box
 
 # Create the chat interface using Gradio Blocks
 with gr.Blocks() as demo:
@@ -185,10 +188,10 @@ 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", value=prompts[0] if prompts else "")
+    system_prompt_dropdown = gr.Dropdown(choices=system_prompts, label="Select System Prompt", value=system_prompts[0] if system_prompts else "")
 
     send_story_button = gr.Button("Send Story")
-    selected_story_textbox = gr.Textbox(lines=10, label="Selected Story", interactive=False)
+    selected_story_display = gr.Textbox(label="Selected Story", interactive=False)
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -204,12 +207,12 @@ with gr.Blocks() as demo:
             comment_input = gr.Textbox(placeholder="Add a comment...", label="Comment")
             save_button = gr.Button("Save Score and Comment")
 
-    data_table = gr.DataFrame(headers=["Timestamp", "User Name", "Model Name", "System Prompt", "Story Name", "User Input", "Chat Response", "Score", "Comment"], visible=True)
+    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_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_display])
     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], outputs=[data_table, comment_input])
+    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()