Merge branch 'fix/query_schedule' of github.com:CaralHsi/MemOSRealPublic into fix/query_schedule

Author: CaralHsi

evaluation/README.md

@@ -1,6 +1,6 @@
 # Evaluation Memory Framework
 
-This repository provides tools and scripts for evaluating the LoCoMo dataset using various models and APIs.
+This repository provides tools and scripts for evaluating the `LoCoMo`, `LongMemEval`, `PrefEval`, `personaMem` dataset using various models and APIs.
 
 ## Installation
 
@@ -68,7 +68,8 @@ First prepare the dataset `longmemeval_s` from https://huggingface.co/datasets/x
 ```
 
 ### PrefEval Evaluation
-To evaluate the **Prefeval** dataset using one of the supported memory frameworks — run the following [script](./scripts/run_prefeval_eval.sh):
+Downloading benchmark_dataset/filtered_inter_turns.json from https://github.com/amazon-science/PrefEval/blob/main/benchmark_dataset/filtered_inter_turns.json and save it as `./data/prefeval/filtered_inter_turns.json`.
+To evaluate the **Prefeval** dataset — run the following [script](./scripts/run_prefeval_eval.sh):
 
 ```bash
 # Edit the configuration in ./scripts/run_prefeval_eval.sh
@@ -83,4 +84,4 @@ get `questions_32k.csv` and `shared_contexts_32k.jsonl` from https://huggingface
 # Specify the model and memory backend you want to use (e.g., mem0, zep, etc.)
 # If you want to use MIRIX, edit the the configuration in ./scripts/personamem/config.yaml
 ./scripts/run_pm_eval.sh
-```
+```

evaluation/scripts/PrefEval/pref_eval.py

@@ -392,23 +392,39 @@ async def main(concurrency_limit: int, input_file: str, output_file: str, output
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Evaluate assistant responses from a JSONL file.")
 
-    parser.add_argument(
-        "--input", type=str, required=True, help="Path to the input JSONL file from pref_memos.py."
-    )
+    parser.add_argument("--input", type=str, required=True, help="Path to the input JSONL file.")
 
     parser.add_argument(
         "--concurrency-limit",
         type=int,
         default=10,
         help="The maximum number of concurrent API calls.",
     )
+
+    parser.add_argument(
+        "--lib",
+        type=str,
+        choices=[
+            "memos-api-online",
+            "mem0",
+            "mem0_graph",
+            "memos-api",
+            "memobase",
+            "memu",
+            "supermemory",
+            "zep",
+        ],
+        default="memos-api",
+        help="Which library to use (used in 'add' mode).",
+    )
+
     args = parser.parse_args()
 
     input_path = args.input
     output_dir = os.path.dirname(input_path)
 
-    output_jsonl_path = os.path.join(output_dir, "eval_pref_memos.jsonl")
-    output_excel_path = os.path.join(output_dir, "eval_pref_memos_summary.xlsx")
+    output_jsonl_path = os.path.join(output_dir, f"eval_pref_{args.lib}.jsonl")
+    output_excel_path = os.path.join(output_dir, f"eval_pref_{args.lib}_summary.xlsx")
 
     asyncio.run(
         main(

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.
     """
@@ -98,7 +103,7 @@ def search_memory_for_line(line_data: tuple, mem_client, top_k_value: int) -> di
                 f"- {entry.get('memory', '')}"
                 for entry in relevant_memories["text_mem"][0]["memories"]
             )
-            + f"\n{relevant_memories['pref_string']}"
+            + f"\n{relevant_memories.get('pref_string', '')}"
         )
 
         memory_tokens_used = len(tokenizer.encode(memories_str))
@@ -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)
             ]