Fix logging

app.py

@@ -31,8 +31,9 @@ model_responses = {"Model_A": "", "Model_B": "", "Model_C": ""}
 
 
 class ModelManager:
-    def __init__(self):
+    def __init__(self, books):
         self.book_model_assignments = {}
+        self.assign_models_to_books(books)
         self.model_A = Client("mep296/llama-3-8b-rephrase-quality")
         self.model_B = Client("mep296/llama-3-8b-rephrase-summarize-quality")
         self.model_C = Client("mep296/llama-3-8b-entigraph-quality")
@@ -83,6 +84,9 @@ class ModelManager:
             }
         return self.book_model_assignments[book]
 
+    def assign_models_to_books(self, books):
+        for book in books:
+            self.get_book_model_mapping(book)
 
 class SheetManager:
     def __init__(self):
@@ -107,9 +111,10 @@ class SheetManager:
 
 class ModelComparisonApp:
     def __init__(self):
-        self.model_manager = ModelManager()
+        self.model_manager = ModelManager(BOOK_MAPPING.keys())
         self.sheet_manager = SheetManager()
         self.votes = []
+        self.selected_book = "Spaceman on a Spree"
         self.selected_book_string = BOOK_MAPPING["Spaceman on a Spree"]
         self.selected_category_string = ""
         self.chat_history_A = []
@@ -171,12 +176,14 @@ class ModelComparisonApp:
             'send': gr.Button("Send", scale=0, variant="primary", interactive=False)
         }
     
-    def respond(self, message, chat_A, chat_B, chat_C):
+    def respond(self, message):
         if not message.strip():
             raise gr.Error("Message cannot be empty!")
             
         prompt = f"{self.selected_book_string} {message}"
-        
+        mapping_dict = self.model_manager.book_model_assignments[self.selected_book]
+        model_order = ["rephrase", "rephrase_summarize", "entigraph"]
+        model_to_index = {model: i for i, model in enumerate(model_order)}
         responses = self.model_manager.get_all_model_responses(prompt)
         
         chats = []
@@ -185,8 +192,9 @@ class ModelComparisonApp:
             chat.append({"role": "user", "content": prompt})
             chat.append({"role": "assistant", "content": response})
             chats.append(chat)
-            
-        return chats
+        
+        reordered_chats = [chats[model_to_index[mapping_dict[model]]] for model in ["Model A", "Model B", "Model C"]]
+        return reordered_chats
     
     def get_votes(self, book, category, question, rating_A, rating_B, rating_C):
         model_mapping = self.model_manager.get_book_model_mapping(book)
@@ -200,12 +208,9 @@ class ModelComparisonApp:
                                  rating_B if model_mapping["Model B"] == "rephrase_summarize" else rating_C,
             "entigraph_rating": rating_A if model_mapping["Model A"] == "entigraph" else 
                         rating_B if model_mapping["Model B"] == "entigraph" else rating_C,
-            "rephrase_response": model_responses["Model_A"] if model_mapping["Model A"] == "rephrase" else 
-                        model_responses["Model_B"] if model_mapping["Model B"] == "rephrase" else model_responses["Model_C"],
-            "rephrase_summarize_response": model_responses["Model_A"] if model_mapping["Model A"] == "rephrase_summarize" else 
-                        model_responses["Model_B"] if model_mapping["Model B"] == "rephrase_summarize" else model_responses["Model_C"],
-            "entigraph_response": model_responses["Model_A"] if model_mapping["Model A"] == "entigraph" else 
-                        model_responses["Model_B"] if model_mapping["Model B"] == "entigraph" else model_responses["Model_C"]
+            "rephrase_response": model_responses["Model_A"],
+            "rephrase_summarize_response": model_responses["Model_B"],
+            "entigraph_response": model_responses["Model_C"]
         }
         
         self.votes.append(rating_data)
@@ -220,6 +225,7 @@ class ModelComparisonApp:
             return gr.update(interactive=bool(model_A and model_B and model_C))
         
         def update_selected_book(book_selection):
+            self.selected_book = book_selection
             self.selected_book_string = BOOK_MAPPING.get(book_selection, "")
             return self.selected_book_string
         
@@ -245,7 +251,7 @@ class ModelComparisonApp:
         
         input_elements['send'].click(
             self.respond,
-            inputs=[input_elements['question']] + list(chat_interfaces.values()),
+            inputs=[input_elements['question']],
             outputs=list(chat_interfaces.values())
         )