Formatting (#2750)

evaluation/src/langmem.py

@@ -28,14 +28,12 @@ def get_answer(question, speaker_1_user_id, speaker_1_memories, speaker_2_user_i
         speaker_1_user_id=speaker_1_user_id,
         speaker_1_memories=speaker_1_memories,
         speaker_2_user_id=speaker_2_user_id,
-        speaker_2_memories=speaker_2_memories
+        speaker_2_memories=speaker_2_memories,
     )
 
     t1 = time.time()
     response = client.chat.completions.create(
-        model=os.getenv("MODEL"),
-        messages=[{"role": "system", "content": prompt}],
-        temperature=0.0
+        model=os.getenv("MODEL"), messages=[{"role": "system", "content": prompt}], temperature=0.0
     )
     t2 = time.time()
     return response.choices[0].message.content, t2 - t1
@@ -59,7 +57,9 @@ def prompt(state):
 
 
 class LangMem:
-    def __init__(self,):
+    def __init__(
+        self,
+    ):
         self.store = InMemoryStore(
             index={
                 "dims": 1536,
@@ -80,18 +80,12 @@ class LangMem:
         )
 
     def add_memory(self, message, config):
-        return self.agent.invoke(
-            {"messages": [{"role": "user", "content": message}]},
-            config=config
-        )
+        return self.agent.invoke({"messages": [{"role": "user", "content": message}]}, config=config)
 
     def search_memory(self, query, config):
         try:
             t1 = time.time()
-            response = self.agent.invoke(
-                {"messages": [{"role": "user", "content": query}]},
-                config=config
-            )
+            response = self.agent.invoke({"messages": [{"role": "user", "content": query}]}, config=config)
             t2 = time.time()
             return response["messages"][-1].content, t2 - t1
         except Exception as e:
@@ -102,7 +96,7 @@ class LangMem:
 class LangMemManager:
     def __init__(self, dataset_path):
         self.dataset_path = dataset_path
-        with open(self.dataset_path, 'r') as f:
+        with open(self.dataset_path, "r") as f:
             self.data = json.load(f)
 
     def process_all_conversations(self, output_file_path):
@@ -123,7 +117,7 @@ class LangMemManager:
 
             # Identify speakers
             for conv in chat_history:
-                speakers.add(conv['speaker'])
+                speakers.add(conv["speaker"])
 
             if len(speakers) != 2:
                 raise ValueError(f"Expected 2 speakers, got {len(speakers)}")
@@ -134,50 +128,52 @@ class LangMemManager:
             # Add memories for each message
             for conv in tqdm(chat_history, desc=f"Processing messages {key}", leave=False):
                 message = f"{conv['timestamp']} | {conv['speaker']}: {conv['text']}"
-                if conv['speaker'] == speaker1:
+                if conv["speaker"] == speaker1:
                     agent1.add_memory(message, config)
-                elif conv['speaker'] == speaker2:
+                elif conv["speaker"] == speaker2:
                     agent2.add_memory(message, config)
                 else:
                     raise ValueError(f"Expected speaker1 or speaker2, got {conv['speaker']}")
 
             # Process questions
             for q in tqdm(questions, desc=f"Processing questions {key}", leave=False):
-                category = q['category']
+                category = q["category"]
 
                 if int(category) == 5:
                     continue
 
-                answer = q['answer']
-                question = q['question']
+                answer = q["answer"]
+                question = q["question"]
                 response1, speaker1_memory_time = agent1.search_memory(question, config)
                 response2, speaker2_memory_time = agent2.search_memory(question, config)
 
-                generated_answer, response_time = get_answer(
-                    question, speaker1, response1, speaker2, response2
-                )
+                generated_answer, response_time = get_answer(question, speaker1, response1, speaker2, response2)
 
-                result[key].append({
-                    "question": question,
-                    "answer": answer,
-                    "response1": response1,
-                    "response2": response2,
-                    "category": category,
-                    "speaker1_memory_time": speaker1_memory_time,
-                    "speaker2_memory_time": speaker2_memory_time,
-                    "response_time": response_time,
-                    'response': generated_answer
-                })
+                result[key].append(
+                    {
+                        "question": question,
+                        "answer": answer,
+                        "response1": response1,
+                        "response2": response2,
+                        "category": category,
+                        "speaker1_memory_time": speaker1_memory_time,
+                        "speaker2_memory_time": speaker2_memory_time,
+                        "response_time": response_time,
+                        "response": generated_answer,
+                    }
+                )
 
             return result
 
         # Use multiprocessing to process conversations in parallel
         with mp.Pool(processes=10) as pool:
-            results = list(tqdm(
-                pool.imap(process_conversation, list(self.data.items())),
-                total=len(self.data),
-                desc="Processing conversations"
-            ))
+            results = list(
+                tqdm(
+                    pool.imap(process_conversation, list(self.data.items())),
+                    total=len(self.data),
+                    desc="Processing conversations",
+                )
+            )
 
         # Combine results from all workers
         for result in results:
@@ -185,5 +181,5 @@ class LangMemManager:
                 OUTPUT[key].extend(items)
 
         # Save final results
-        with open(output_file_path, 'w') as f:
+        with open(output_file_path, "w") as f:
             json.dump(OUTPUT, f, indent=4)