add breakpoint in eval scripts

Author: Nyakult

evaluation/scripts/PrefEval/pref_mem0.py

@@ -29,7 +29,13 @@
 
 
 def add_memory_for_line(
-    line_data: tuple, mem_client, num_irrelevant_turns: int, lib: str, version: str
+    line_data: tuple,
+    mem_client,
+    num_irrelevant_turns: int,
+    lib: str,
+    version: str,
+    success_records,
+    f,
 ) -> dict:
     """
     Adds conversation memory for a single line of data to MemOS and returns the data with a persistent user_id.
@@ -46,13 +52,22 @@ def add_memory_for_line(
         elif num_irrelevant_turns == 300:
             conversation = conversation + irre_300
 
-        turns_add = 5
         start_time_add = time.monotonic()
-        if conversation:
-            for chunk_start in range(0, len(conversation), turns_add * 2):
-                chunk = conversation[chunk_start : chunk_start + turns_add * 2]
-                timestamp_add = int(time.time() * 100)
-                mem_client.add(messages=chunk, user_id=user_id, timestamp=timestamp_add)
+
+        for idx, _ in enumerate(conversation[::2]):
+            msg_idx = idx * 2
+            record_id = f"{lib}_user_pref_eval_{i}_{version}_{str(msg_idx)}"
+            timestamp_add = int(time.time() * 100)
+
+            if record_id not in success_records:
+                mem_client.add(
+                    messages=conversation[msg_idx : msg_idx + 2],
+                    user_id=user_id,
+                    timestamp=timestamp_add,
+                )
+                f.write(f"{record_id}\n")
+                f.flush()
+
         end_time_add = time.monotonic()
         add_duration = end_time_add - start_time_add
 
@@ -210,6 +225,15 @@ def main():
     from utils.client import Mem0Client
 
     mem_client = Mem0Client(enable_graph="graph" in args.lib)
+    os.makedirs(f"results/prefeval/{args.lib}_{args.version}", exist_ok=True)
+    success_records = set()
+    record_file = f"results/prefeval/{args.lib}_{args.version}/success_records.txt"
+    if os.path.exists(record_file):
+        print(f"Loading existing success records from {record_file}...")
+        with open(record_file, encoding="utf-8") as f:
+            for i in f.readlines():
+                success_records.add(i.strip())
+        print(f"Loaded {len(success_records)} records.")
 
     if args.mode == "add":
         print(f"Running in 'add' mode. Ingesting memories from '{args.input}'...")
@@ -218,6 +242,7 @@ def main():
         with (
             open(args.output, "w", encoding="utf-8") as outfile,
             concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor,
+            open(record_file, "a+", encoding="utf-8") as f,
         ):
             futures = [
                 executor.submit(
@@ -227,6 +252,8 @@ def main():
                     args.add_turn,
                     args.lib,
                     args.version,
+                    success_records,
+                    f,
                 )
                 for i, line in enumerate(lines)
             ]

evaluation/scripts/PrefEval/pref_memobase.py

@@ -12,7 +12,6 @@
 from openai import OpenAI
 from tqdm import tqdm
 
-
 ROOT_DIR = os.path.dirname(
     os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 )
@@ -28,16 +27,20 @@
 
 
 def add_memory_for_line(
-    line_data: tuple, mem_client, num_irrelevant_turns: int, lib: str, version: str
+    line_data: tuple,
+    mem_client,
+    num_irrelevant_turns: int,
+    lib: str,
+    version: str,
+    success_records,
+    f,
 ) -> dict:
     """
     Adds conversation memory for a single line of data to MemOS and returns the data with a persistent user_id.
     """
     i, line = line_data
     user_id = f"{lib}_user_pref_eval_{i}_{version}"
     mem_client.delete_user(user_id)
-    user_id = mem_client.client.add_user({"user_id": user_id})
-    print("user_id:", user_id)
     try:
         original_data = json.loads(line)
         conversation = original_data.get("conversation", [])
@@ -63,7 +66,14 @@ def add_memory_for_line(
                         "created_at": timestamp_add,
                     }
                 )
-            mem_client.add(messages=messages, user_id=user_id)
+            for idx, _ in enumerate(conversation[::2]):
+                msg_idx = idx * 2
+                record_id = f"{lib}_user_pref_eval_{i}_{version}_{str(msg_idx)}"
+
+                if record_id not in success_records:
+                    mem_client.add(messages=conversation[msg_idx : msg_idx + 2], user_id=user_id)
+                    f.write(f"{record_id}\n")
+                    f.flush()
 
         end_time_add = time.monotonic()
         add_duration = end_time_add - start_time_add
@@ -222,13 +232,24 @@ def main():
 
     mem_client = MemobaseClient()
 
+    os.makedirs(f"results/prefeval/{args.lib}_{args.version}", exist_ok=True)
+    success_records = set()
+    record_file = f"results/prefeval/{args.lib}_{args.version}/success_records.txt"
+    if os.path.exists(record_file):
+        print(f"Loading existing success records from {record_file}...")
+        with open(record_file, encoding="utf-8") as f:
+            for i in f.readlines():
+                success_records.add(i.strip())
+        print(f"Loaded {len(success_records)} records.")
+
     if args.mode == "add":
         print(f"Running in 'add' mode. Ingesting memories from '{args.input}'...")
         print(f"Adding {args.add_turn} irrelevant turns.")
         print(f"Using {args.max_workers} workers.")
         with (
             open(args.output, "w", encoding="utf-8") as outfile,
             concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor,
+            open(record_file, "a+", encoding="utf-8") as f,
         ):
             futures = [
                 executor.submit(
@@ -238,6 +259,8 @@ def main():
                     args.add_turn,
                     args.lib,
                     args.version,
+                    success_records,
+                    f,
                 )
                 for i, line in enumerate(lines)
             ]

evaluation/scripts/PrefEval/pref_memos.py

@@ -12,7 +12,6 @@
 from openai import OpenAI
 from tqdm import tqdm
 
-
 ROOT_DIR = os.path.dirname(
     os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 )
@@ -21,7 +20,6 @@
 sys.path.insert(0, ROOT_DIR)
 sys.path.insert(0, EVAL_SCRIPTS_DIR)
 
-
 load_dotenv()
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 BASE_URL = os.getenv("OPENAI_BASE_URL")
@@ -30,8 +28,8 @@
 
 
 def add_memory_for_line(
-    line_data: tuple, mem_client, num_irrelevant_turns: int, lib: str, version: str
-) -> dict:
+    line_data, mem_client, num_irrelevant_turns, lib, version, success_records, f
+):
     """
     Adds conversation memory for a single line of data to MemOS and returns the data with a persistent user_id.
     """
@@ -47,15 +45,22 @@ def add_memory_for_line(
         elif num_irrelevant_turns == 300:
             conversation = conversation + irre_300
 
-        turns_add = 5
         start_time_add = time.monotonic()
-        if conversation:
-            if os.getenv("PRE_SPLIT_CHUNK", "false").lower() == "true":
-                for chunk_start in range(0, len(conversation), turns_add * 2):
-                    chunk = conversation[chunk_start : chunk_start + turns_add * 2]
-                    mem_client.add(messages=chunk, user_id=user_id, conv_id=None, batch_size=2)
-            else:
-                mem_client.add(messages=conversation, user_id=user_id, conv_id=None, batch_size=2)
+
+        for idx, _ in enumerate(conversation[::2]):
+            msg_idx = idx * 2
+            record_id = f"{lib}_user_pref_eval_{i}_{version}_{str(msg_idx)}"
+
+            if record_id not in success_records:
+                mem_client.add(
+                    messages=conversation[msg_idx : msg_idx + 2],
+                    user_id=user_id,
+                    conv_id=None,
+                    batch_size=2,
+                )
+                f.write(f"{record_id}\n")
+                f.flush()
+
         end_time_add = time.monotonic()
         add_duration = end_time_add - start_time_add
 
@@ -68,7 +73,7 @@ def add_memory_for_line(
         return None
 
 
-def search_memory_for_line(line_data: tuple, mem_client, top_k_value: int) -> dict:
+def search_memory_for_line(line_data, mem_client, top_k_value):
     """
     Processes a single line of data, searching memory based on the question.
     """
@@ -120,7 +125,7 @@ def search_memory_for_line(line_data: tuple, mem_client, top_k_value: int) -> di
         return None
 
 
-def generate_response_for_line(line_data: tuple, openai_client: OpenAI, lib: str) -> dict:
+def generate_response_for_line(line_data, openai_client, lib):
     """
     Generates a response for a single line of data using pre-fetched memories.
     """
@@ -195,7 +200,7 @@ def main():
     parser.add_argument(
         "--lib",
         type=str,
-        choices=["memos-api", "memos-local"],
+        choices=["memos-api", "memos-api-online"],
         default="memos-api",
         help="Which MemOS library to use (used in 'add' mode).",
     )
@@ -218,9 +223,22 @@ def main():
         print(f"Error: Input file '{args.input}' not found")
         return
 
-    from utils.client import MemosApiClient
+    from utils.client import MemosApiClient, MemosApiOnlineClient
+
+    if args.lib == "memos-api":
+        mem_client = MemosApiClient()
+    elif args.lib == "memos-api-online":
+        mem_client = MemosApiOnlineClient()
 
-    mem_client = MemosApiClient()
+    os.makedirs(f"results/prefeval/{args.lib}_{args.version}", exist_ok=True)
+    success_records = set()
+    record_file = f"results/prefeval/{args.lib}_{args.version}/success_records.txt"
+    if os.path.exists(record_file):
+        print(f"Loading existing success records from {record_file}...")
+        with open(record_file, encoding="utf-8") as f:
+            for i in f.readlines():
+                success_records.add(i.strip())
+        print(f"Loaded {len(success_records)} records.")
 
     if args.mode == "add":
         print(f"Running in 'add' mode. Ingesting memories from '{args.input}'...")
@@ -229,6 +247,7 @@ def main():
         with (
             open(args.output, "w", encoding="utf-8") as outfile,
             concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor,
+            open(record_file, "a+", encoding="utf-8") as record_f,
         ):
             futures = [
                 executor.submit(
@@ -238,6 +257,8 @@ def main():
                     args.add_turn,
                     args.lib,
                     args.version,
+                    success_records,
+                    record_f,
                 )
                 for i, line in enumerate(lines)
             ]

evaluation/scripts/PrefEval/pref_memu.py

@@ -14,7 +14,6 @@
 from openai import OpenAI
 from tqdm import tqdm
 
-
 ROOT_DIR = os.path.dirname(
     os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 )
@@ -30,7 +29,13 @@
 
 
 def add_memory_for_line(
-    line_data: tuple, mem_client, num_irrelevant_turns: int, lib: str, version: str
+    line_data: tuple,
+    mem_client,
+    num_irrelevant_turns: int,
+    lib: str,
+    version: str,
+    success_records,
+    f,
 ) -> dict:
     """
     Adds conversation memory for a single line of data to MemOS and returns the data with a persistent user_id.
@@ -47,19 +52,21 @@ def add_memory_for_line(
         elif num_irrelevant_turns == 300:
             conversation = conversation + irre_300
 
-        turns_add = 5
         start_time_add = time.monotonic()
-        if conversation:
-            if os.getenv("PRE_SPLIT_CHUNK", "false").lower() == "true":
-                for chunk_start in range(0, len(conversation), turns_add * 2):
-                    chunk = conversation[chunk_start : chunk_start + turns_add * 2]
-                    mem_client.add(
-                        messages=chunk, user_id=user_id, iso_date=datetime.now().isoformat()
-                    )
-            else:
+
+        for idx, _ in enumerate(conversation[::2]):
+            msg_idx = idx * 2
+            record_id = f"{lib}_user_pref_eval_{i}_{version}_{str(msg_idx)}"
+
+            if record_id not in success_records:
                 mem_client.add(
-                    messages=conversation, user_id=user_id, iso_date=datetime.now().isoformat()
+                    messages=conversation[msg_idx : msg_idx + 2],
+                    user_id=user_id,
+                    iso_date=datetime.now().isoformat(),
                 )
+                f.write(f"{record_id}\n")
+                f.flush()
+
         end_time_add = time.monotonic()
         add_duration = end_time_add - start_time_add
 
@@ -219,13 +226,24 @@ def main():
 
     mem_client = MemuClient()
 
+    os.makedirs(f"results/prefeval/{args.lib}_{args.version}", exist_ok=True)
+    success_records = set()
+    record_file = f"results/prefeval/{args.lib}_{args.version}/success_records.txt"
+    if os.path.exists(record_file):
+        print(f"Loading existing success records from {record_file}...")
+        with open(record_file, encoding="utf-8") as f:
+            for i in f.readlines():
+                success_records.add(i.strip())
+        print(f"Loaded {len(success_records)} records.")
+
     if args.mode == "add":
         print(f"Running in 'add' mode. Ingesting memories from '{args.input}'...")
         print(f"Adding {args.add_turn} irrelevant turns.")
         print(f"Using {args.max_workers} workers.")
         with (
             open(args.output, "w", encoding="utf-8") as outfile,
             concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor,
+            open(record_file, "a+", encoding="utf-8") as f,
         ):
             futures = [
                 executor.submit(
@@ -235,6 +253,8 @@ def main():
                     args.add_turn,
                     args.lib,
                     args.version,
+                    success_records,
+                    f,
                 )
                 for i, line in enumerate(lines)
             ]