Feat/pref optimize update (#409)

* add hybrid search and fine extractor

* add dialog and modify spliter chunk

* optmize the update and retriever code

* modify pref field

* add pref mem update srategy

* add pref mem update srategy

* fix bug in pre_commit

---------

Co-authored-by: yuan.wang <[email protected]>

Author: Wang-Daoji

docker/requirements.txt

@@ -157,4 +157,4 @@ volcengine-python-sdk==4.0.6
 watchfiles==1.1.0
 websockets==15.0.1
 xlrd==2.0.2
-xlsxwriter==3.2.5
+xlsxwriter==3.2.5

evaluation/scripts/PrefEval/pref_memos.py

@@ -53,9 +53,9 @@ def add_memory_for_line(
             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)
+                    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)
+                mem_client.add(messages=conversation, user_id=user_id, conv_id=None, batch_size=2)
         end_time_add = time.monotonic()
         add_duration = end_time_add - start_time_add
 
@@ -98,7 +98,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_mem']}"
+            + f"\n{relevant_memories['pref_string']}"
         )
 
         memory_tokens_used = len(tokenizer.encode(memories_str))

evaluation/scripts/locomo/locomo_search.py

@@ -107,11 +107,11 @@ def memos_api_search(
 
     speaker_a_context = (
         "\n".join([i["memory"] for i in search_a_results["text_mem"][0]["memories"]])
-        + f"\n{search_a_results['pref_mem']}"
+        + f"\n{search_a_results['pref_string']}"
     )
     speaker_b_context = (
         "\n".join([i["memory"] for i in search_b_results["text_mem"][0]["memories"]])
-        + f"\n{search_b_results['pref_mem']}"
+        + f"\n{search_b_results['pref_string']}"
     )
 
     context = TEMPLATE_MEMOS.format(

evaluation/scripts/longmemeval/lme_search.py

@@ -46,7 +46,7 @@ def memos_search(client, query, user_id, top_k):
     results = client.search(query=query, user_id=user_id, top_k=top_k)
     context = (
         "\n".join([i["memory"] for i in results["text_mem"][0]["memories"]])
-        + f"\n{results['pref_mem']}"
+        + f"\n{results['pref_string']}"
     )
     context = MEMOS_CONTEXT_TEMPLATE.format(user_id=user_id, memories=context)
     duration_ms = (time() - start) * 1000

evaluation/scripts/personamem/pm_ingestion.py

@@ -31,10 +31,10 @@ def ingest_session(session, user_id, session_id, frame, client):
         if os.getenv("PRE_SPLIT_CHUNK") == "true":
             for i in range(0, len(session), 10):
                 messages = session[i : i + 10]
-                client.add(messages=messages, user_id=user_id, conv_id=session_id)
+                client.add(messages=messages, user_id=user_id, conv_id=session_id, batch_size=2)
                 print(f"[{frame}] ✅ Session [{session_id}]: Ingested {len(messages)} messages")
         else:
-            client.add(messages=session, user_id=user_id, conv_id=session_id)
+            client.add(messages=session, user_id=user_id, conv_id=session_id, batch_size=2)
             print(f"[{frame}] ✅ Session [{session_id}]: Ingested {len(session)} messages")
     elif frame == "memobase":
         for _idx, msg in enumerate(session):

evaluation/scripts/personamem/pm_search.py

@@ -84,7 +84,7 @@ def memos_search(client, user_id, query, top_k):
     results = client.search(query=query, user_id=user_id, top_k=top_k)
     search_memories = (
         "\n".join(item["memory"] for cube in results["text_mem"] for item in cube["memories"])
-        + f"\n{results['pref_mem']}"
+        + f"\n{results['pref_string']}"
     )
     context = MEMOS_CONTEXT_TEMPLATE.format(user_id=user_id, memories=search_memories)
 

evaluation/scripts/utils/client.py

@@ -182,6 +182,7 @@ def search(self, query, user_id, top_k):
                 "conversation_id": "",
                 "top_k": top_k,
                 "mode": "mixture",
+                "handle_pref_mem": False,
             },
             ensure_ascii=False,
         )

src/memos/api/routers/server_router.py

@@ -303,18 +303,15 @@ def _post_process_pref_mem(
     mem_cube_id: str,
     handle_pref_mem: bool,
 ):
-    if os.getenv("RETURN_ORIGINAL_PREF_MEM", "false").lower() == "true" and pref_formatted_mem:
-        memories_result["prefs"] = []
-        memories_result["prefs"].append(
+    if handle_pref_mem:
+        memories_result["pref_mem"].append(
             {
                 "cube_id": mem_cube_id,
                 "memories": pref_formatted_mem,
             }
         )
-
-    if handle_pref_mem:
         pref_instruction: str = instruct_completion(pref_formatted_mem)
-        memories_result["pref_mem"] = pref_instruction
+        memories_result["pref_string"] = pref_instruction
 
     return memories_result
 
@@ -333,7 +330,8 @@ def search_memories(search_req: APISearchRequest):
         "text_mem": [],
         "act_mem": [],
         "para_mem": [],
-        "pref_mem": "",
+        "pref_mem": [],
+        "pref_string": "",
     }
 
     search_mode = search_req.mode

src/memos/memories/textual/item.py

@@ -194,7 +194,7 @@ class PreferenceTextualMemoryMetadata(TextualMemoryMetadata):
         default="explicit_preference", description="Type of preference."
     )
     dialog_id: str | None = Field(default=None, description="ID of the dialog.")
-    dialog_str: str | None = Field(default=None, description="String of the dialog.")
+    original_text: str | None = Field(default=None, description="String of the dialog.")
     embedding: list[float] | None = Field(default=None, description="Vector of the dialog.")
     explicit_preference: str | None = Field(default=None, description="Explicit preference.")
     created_at: str | None = Field(default=None, description="Timestamp of the dialog.")