format code

Author: Nyakult

evaluation/scripts/locomo/locomo_ingestion.py

@@ -1,20 +1,18 @@
-import asyncio
 import os
 import sys
+import argparse
+import concurrent.futures
+import time
+from datetime import datetime, timezone
+import pandas as pd
+from dotenv import load_dotenv
 
 ROOT_DIR = os.path.dirname(
     os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 )
 EVAL_SCRIPTS_DIR = os.path.join(ROOT_DIR, "evaluation", "scripts")
 sys.path.insert(0, ROOT_DIR)
 sys.path.insert(0, EVAL_SCRIPTS_DIR)
-
-import argparse
-import concurrent.futures
-import time
-from datetime import datetime, timezone
-import pandas as pd
-from dotenv import load_dotenv
 from prompts import custom_instructions
 
 
@@ -64,20 +62,13 @@ def ingest_session(client, session, frame, version, metadata):
         client.add(speaker_a_messages, speaker_a_user_id)
         client.add(speaker_b_messages, speaker_b_user_id)
     elif frame == "memu":
-        # speaker_a_user_id = metadata['speaker_a']
-        # speaker_b_user_id = metadata['speaker_b']
-        # client.agent_id = speaker_b_user_id
         client.add(speaker_a_messages, speaker_a_user_id, iso_date)
-        # client.agent_id = speaker_a_user_id
         client.add(speaker_b_messages, speaker_b_user_id, iso_date)
     elif frame == "supermemory":
         for m in speaker_a_messages:
             m["chat_time"] = iso_date
         for m in speaker_b_messages:
             m["chat_time"] = iso_date
-        # seems like user_id can not be too long
-        speaker_a_user_id = f"lcm{conv_idx}a_{version}"
-        speaker_b_user_id = f"lcm{conv_idx}b_{version}"
         client.add(speaker_a_messages, speaker_a_user_id)
         client.add(speaker_b_messages, speaker_b_user_id)
 
@@ -114,11 +105,8 @@ def process_user(conv_idx, frame, locomo_df, version):
         client = memobase_client()
         all_users = client.client.get_all_users(limit=5000)
         for user in all_users:
-            try:
-                if user["additional_fields"]["user_id"] in [speaker_a_user_id, speaker_b_user_id]:
-                    client.client.delete_user(user["id"])
-            except:
-                pass
+            if user["additional_fields"]["user_id"] in [speaker_a_user_id, speaker_b_user_id]:
+                client.client.delete_user(user["id"])
         speaker_a_user_id = client.client.add_user({"user_id": speaker_a_user_id})
         speaker_b_user_id = client.client.add_user({"user_id": speaker_b_user_id})
     elif frame == "memu":
@@ -129,7 +117,6 @@ def process_user(conv_idx, frame, locomo_df, version):
         from utils.client import supermemory_client
 
         client = supermemory_client()
-
     sessions_to_process = []
     for session_idx in range(max_session_count):
         session_key = f"session_{session_idx}"

evaluation/scripts/locomo/locomo_responses.py

@@ -112,8 +112,6 @@ async def main(frame, version="default"):
 
     os.makedirs("data", exist_ok=True)
 
-    # print(all_responses)
-
     with open(response_path, "w") as f:
         json.dump(all_responses, f, indent=2)
         print("Save response results")

evaluation/scripts/locomo/locomo_search.py

@@ -1,6 +1,13 @@
 import os
 import sys
-
+import argparse
+import json
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from time import time
+import pandas as pd
+from dotenv import load_dotenv
+from tqdm import tqdm
 
 ROOT_DIR = os.path.dirname(
     os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -9,16 +16,8 @@
 
 sys.path.insert(0, ROOT_DIR)
 sys.path.insert(0, EVAL_SCRIPTS_DIR)
-from prompts import TEMPLATE_MEM0, TEMPLATE_MEMOBASE, TEMPLATE_MEMOS, TEMPLATE_MEM0_GRAPH
 
-import argparse
-import json
-from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from time import time
-import pandas as pd
-from dotenv import load_dotenv
-from tqdm import tqdm
+from prompts import TEMPLATE_MEM0, TEMPLATE_MEMOBASE, TEMPLATE_MEMOS, TEMPLATE_MEM0_GRAPH
 
 
 def mem0_search(client, query, speaker_a_user_id, speaker_b_user_id, top_k, speaker_a, speaker_b):
@@ -168,7 +167,7 @@ def supermemory_search(
 
 
 def search_query(client, query, metadata, frame, version, top_k=20):
-    conv_id = metadata.get("conv_id")
+    _conv_id = metadata.get("conv_id")
     speaker_a = metadata.get("speaker_a")
     speaker_b = metadata.get("speaker_b")
     speaker_a_user_id = metadata.get("speaker_a_user_id")
@@ -247,13 +246,10 @@ def process_user(conv_idx, locomo_df, frame, version, top_k=20, num_workers=1):
         client = memobase_client()
         users = client.client.get_all_users(limit=5000)
         for u in users:
-            try:
-                if u["additional_fields"]["user_id"] == speaker_a_user_id:
-                    speaker_a_user_id = u["id"]
-                if u["additional_fields"]["user_id"] == speaker_b_user_id:
-                    speaker_b_user_id = u["id"]
-            except:
-                pass
+            if u["additional_fields"]["user_id"] == speaker_a_user_id:
+                speaker_a_user_id = u["id"]
+            if u["additional_fields"]["user_id"] == speaker_b_user_id:
+                speaker_b_user_id = u["id"]
     elif frame == "memu":
         from utils.client import memu_client
 

evaluation/scripts/longmemeval/lme_ingestion.py

@@ -1,14 +1,11 @@
 import argparse
 import os
 import sys
-
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime, timezone
-
 import pandas as pd
-
 from tqdm import tqdm
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 
 def ingest_session(session, date, user_id, session_id, frame, client):
@@ -18,7 +15,7 @@ def ingest_session(session, date, user_id, session_id, frame, client):
             messages.append({"role": msg["role"], "content": msg["content"][:8000]})
             client.add(messages, user_id, int(date.timestamp()))
     elif frame == "memobase":
-        for idx, msg in enumerate(session):
+        for _idx, msg in enumerate(session):
             messages.append(
                 {
                     "role": msg["role"],
@@ -39,11 +36,11 @@ def ingest_session(session, date, user_id, session_id, frame, client):
         if messages:
             client.add(messages=messages, user_id=user_id, conv_id=session_id)
     elif frame == "memu":
-        for idx, msg in enumerate(session):
+        for _idx, msg in enumerate(session):
             messages.append({"role": msg["role"], "content": msg["content"][:8000]})
         client.add(messages, user_id, date.isoformat())
     elif frame == "supermemory":
-        for idx, msg in enumerate(session):
+        for _idx, msg in enumerate(session):
             messages.append(
                 {
                     "role": msg["role"],
@@ -84,11 +81,8 @@ def ingest_conv(lme_df, version, conv_idx, frame, success_records, f):
         client = memobase_client()
         all_users = client.client.get_all_users(limit=5000)
         for user in all_users:
-            try:
-                if user["additional_fields"]["user_id"] == user_id:
-                    client.client.delete_user(user["id"])
-            except:
-                pass
+            if user["additional_fields"]["user_id"] == user_id:
+                client.client.delete_user(user["id"])
         user_id = client.client.add_user({"user_id": user_id})
     elif frame == "memu":
         from utils.client import memu_client
@@ -135,8 +129,9 @@ def main(frame, version, num_workers=2):
     success_records = []
     record_file = f"results/lme/{frame}-{version}/success_records.txt"
     if os.path.exists(record_file):
-        for i in open(record_file, "r").readlines():
-            success_records.append(i.strip())
+        with open(record_file, "r") as f:
+            for i in f.readlines():
+                success_records.append(i.strip())
 
     f = open(record_file, "a+")
 

evaluation/scripts/longmemeval/lme_search.py

@@ -116,11 +116,8 @@ def process_user(lme_df, conv_idx, frame, version, top_k=20):
         client = memobase_client()
         users = client.client.get_all_users(limit=5000)
         for u in users:
-            try:
-                if u["additional_fields"]["user_id"] == user_id:
-                    user_id = u["id"]
-            except:
-                pass
+            if u["additional_fields"]["user_id"] == user_id:
+                user_id = u["id"]
         context, duration_ms = memobase_search(client, question, user_id, top_k)
     elif frame == "memos-api":
         from utils.client import memos_api_client
@@ -196,7 +193,7 @@ def main(frame, version, top_k=20, num_workers=2):
         for future in tqdm(
             as_completed(future_to_idx), total=num_multi_sessions, desc="📊 Processing users"
         ):
-            idx = future_to_idx[future]
+            _idx = future_to_idx[future]
             search_results = future.result()
             for user_id, results in search_results.items():
                 all_search_results[user_id].extend(results)