fix bugs in text mem with neo4j backend, and set huggingface backend to be singleton

Author: tangg555

evaluation/scripts/temporal_locomo/locomo_processor.py

@@ -100,7 +100,8 @@ def process_user(self, conv_id, locomo_df, frame, version, top_k=20):
 
         oai_client = OpenAI(api_key=self.openai_api_key, base_url=self.openai_base_url)
 
-        self.pre_context_cache[conv_id] = None
+        with self.stats_lock:
+            self.pre_context_cache[conv_id] = None
 
         def process_qa(qa):
             try:
@@ -119,12 +120,14 @@ def process_qa(qa):
                     context = ""
 
                 # ==== Context Answerability Analysis (for memos_scheduler only) ====
-
+                can_answer = False
+                can_answer_duration_ms = 0.0
                 if self.pre_context_cache[conv_id] is not None:
+                    can_answer_start = time()
                     can_answer = self.analyze_context_answerability(
                         self.pre_context_cache[conv_id], query, oai_client
                     )
-
+                    can_answer_duration_ms = (time() - can_answer_start) * 1000
                     # Update statistics
                     with self.stats_lock:
                         self.stats[self.frame][self.version]["memory_stats"]["total_queries"] += 1
@@ -151,8 +154,8 @@ def process_qa(qa):
                             hit_rate
                         )
                         self.save_stats()
-
-                self.pre_context_cache[conv_id] = context
+                with self.stats_lock:
+                    self.pre_context_cache[conv_id] = context
 
                 self.print_eval_info()
 
@@ -174,6 +177,7 @@ def process_qa(qa):
                     "search_context": context,
                     "response_duration_ms": response_duration_ms,
                     "search_duration_ms": search_duration_ms,
+                    "can_answer_duration_ms": can_answer_duration_ms,
                     "can_answer": can_answer if frame == "memos_scheduler" else None,
                 }
             except Exception as e:
@@ -191,16 +195,19 @@ def process_qa(qa):
                         if result["search_context"]
                         else "No context"
                     )
-                    print(
-                        {
-                            "question": result["question"][:100],
-                            "answer": result["answer"][:100],
-                            "category": result["category"],
-                            "golden_answer": result["golden_answer"],
-                            "search_context": context_preview[:100],
-                            "search_duration_ms": result["search_duration_ms"],
-                        }
-                    )
+                    if "can_answer" in result:
+                        print("Print can_answer examples")
+                        print(
+                            {
+                                "question": result["question"][:100],
+                                "pre context can answer": result["can_answer"],
+                                "answer": result["answer"][:100],
+                                "category": result["category"],
+                                "golden_answer": result["golden_answer"],
+                                "search_context": context_preview[:100],
+                                "search_duration_ms": result["search_duration_ms"],
+                            }
+                        )
 
                     search_results[conv_id].append(
                         {

evaluation/scripts/temporal_locomo/modules/client_manager.py

@@ -153,8 +153,6 @@ def get_client_from_storage(
                 # Replace the original scheduler
                 mos.mem_scheduler = scheduler_for_eval
 
-                # stop mem_scheduler thread
-                mos.mem_scheduler.stop()
             return mos
 
     def locomo_response(self, frame, llm_client, context: str, question: str) -> str:

examples/data/config/mem_scheduler/mem_cube_config.yaml

@@ -34,7 +34,7 @@ act_mem:
   config:
     memory_filename: "activation_memory.pickle"
     extractor_llm:
-      backend: "huggingface"
+      backend: "huggingface_singleton"
       config:
         model_name_or_path: "Qwen/Qwen3-1.7B"
         temperature: 0.8
@@ -48,7 +48,7 @@ para_mem:
   config:
     memory_filename: "parametric_memory.adapter"
     extractor_llm:
-      backend: "huggingface"
+      backend: "huggingface_singleton"
       config:
         model_name_or_path: "Qwen/Qwen3-1.7B"
         temperature: 0.8

examples/data/config/mem_scheduler/memos_config_w_optimized_scheduler_and_openai.yaml

@@ -1,6 +1,6 @@
 user_id: "root"
 chat_model:
-  backend: "huggingface"
+  backend: "huggingface_singleton"
   config:
     model_name_or_path: "Qwen/Qwen3-1.7B"
     temperature: 0.1

examples/data/config/mem_scheduler/memos_config_w_scheduler.yaml

@@ -1,6 +1,6 @@
 user_id: "root"
 chat_model:
-  backend: "huggingface"
+  backend: "huggingface_singleton"
   config:
     model_name_or_path: "Qwen/Qwen3-1.7B"
     temperature: 0.1

examples/data/config/mem_scheduler/memos_config_w_scheduler_and_openai.yaml

@@ -1,6 +1,6 @@
 user_id: "root"
 chat_model:
-  backend: "huggingface"
+  backend: "huggingface_singleton"
   config:
     model_name_or_path: "Qwen/Qwen3-1.7B"
     temperature: 0.1

examples/mem_scheduler/memos_w_scheduler_for_test.py

@@ -186,7 +186,7 @@ def init_task():
         )
 
     # Initialization
-    print("🔧 Initializing MOS with Enhanced Scheduler...")
+    print("🔧 Initializing MOS with Scheduler...")
     mos = MOSForTestScheduler(mos_config)
 
     user_id = "user_1"

src/memos/graph_dbs/neo4j.py

@@ -1,3 +1,4 @@
+import json
 import time
 
 from datetime import datetime
@@ -174,6 +175,12 @@ def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None:
                 n.updated_at = datetime($updated_at),
                 n += $metadata
         """
+
+        # serialization
+        if metadata["sources"]:
+            for idx in range(len(metadata["sources"])):
+                metadata["sources"][idx] = json.dumps(metadata["sources"][idx])
+
         with self.driver.session(database=self.db_name) as session:
             session.run(
                 query,
@@ -1128,4 +1135,14 @@ def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]:
                 node[time_field] = node[time_field].isoformat()
         node.pop("user_name", None)
 
+        # serialization
+        if node["sources"]:
+            for idx in range(len(node["sources"])):
+                if not (
+                    isinstance(node["sources"][idx], str)
+                    and node["sources"][idx][0] == "{"
+                    and node["sources"][idx][0] == "}"
+                ):
+                    break
+                node["sources"][idx] = json.loads(node["sources"][idx])
         return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node}