fix the bugs in rule-based baselines, and change the temporal data sorting strategy

Author: tangg555

evaluation/scripts/temporal_locomo/locomo_processor.py

@@ -1,10 +1,8 @@
 import json
 import sys
-import traceback
 
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from datetime import datetime
 from pathlib import Path
 from time import time
 
@@ -20,10 +18,8 @@
     SEARCH_PROMPT_MEMOS,
     SEARCH_PROMPT_ZEP,
 )
-from modules.schemas import RecordingCase
+from modules.schemas import ContextUpdateMethod, RecordingCase
 from modules.utils import save_evaluation_cases
-from openai import OpenAI
-from tqdm import tqdm
 
 from memos.log import get_logger
 
@@ -57,68 +53,24 @@ def __init__(self, args):
 
         self.processed_data_dir = self.result_dir / "processed_data"
 
-    # -------------------------------
-    # Refactor helpers for process_user
-    # -------------------------------
-
-    def _initialize_conv_stats(self):
-        """Create a fresh statistics dictionary for a conversation."""
-        return {
-            "total_queries": 0,
-            "can_answer_count": 0,
-            "cannot_answer_count": 0,
-            "answer_hit_rate": 0.0,
-            "response_failure": 0,
-            "response_count": 0,
-        }
-
-    def _build_day_groups(self, temporal_conv):
-        """Build mapping day_id -> qa_pairs from a temporal conversation dict."""
-        day_groups = {}
-        for day_id, day_data in temporal_conv.get("days", {}).items():
-            day_groups[day_id] = day_data.get("qa_pairs", [])
-        return day_groups
-
-    def _build_metadata(self, speaker_a, speaker_b, speaker_a_user_id, speaker_b_user_id, conv_id):
-        """Assemble metadata for downstream calls."""
-        return {
-            "speaker_a": speaker_a,
-            "speaker_b": speaker_b,
-            "speaker_a_user_id": speaker_a_user_id,
-            "speaker_b_user_id": speaker_b_user_id,
-            "conv_id": conv_id,
-        }
-
-    def _get_clients(self, frame, speaker_a_user_id, speaker_b_user_id, conv_id, version, top_k):
-        """Return (client, reversed_client) according to the target frame."""
-        reversed_client = None
-        if frame in [MEMOS_MODEL, MEMOS_SCHEDULER_MODEL]:
-            client = self.get_client_from_storage(frame, speaker_a_user_id, version, top_k=top_k)
-            reversed_client = self.get_client_from_storage(
-                frame, speaker_b_user_id, version, top_k=top_k
-            )
+    def update_context(self, conv_id, method, **kwargs):
+        if method == ContextUpdateMethod.DIRECT:
+            if "cur_context" not in kwargs:
+                raise ValueError("cur_context is required for DIRECT update method")
+            cur_context = kwargs["cur_context"]
+            self.pre_context_cache[conv_id] = cur_context
+        elif method == ContextUpdateMethod.TEMPLATE:
+            if "query" not in kwargs or "answer" not in kwargs:
+                raise ValueError("query and answer are required for TEMPLATE update method")
+            self._update_context_template(conv_id, kwargs["query"], kwargs["answer"])
         else:
-            client = self.get_client_from_storage(frame, conv_id, version)
-        return client, reversed_client
-
-    def _save_conv_stats(self, conv_id, frame, version, conv_stats, conv_stats_path):
-        """Persist per-conversation stats to disk."""
-        conv_stats_data = {
-            "conversation_id": conv_id,
-            "frame": frame,
-            "version": version,
-            "statistics": conv_stats,
-            "timestamp": str(datetime.now()),
-        }
-        with open(conv_stats_path, "w") as fw:
-            json.dump(conv_stats_data, fw, indent=2, ensure_ascii=False)
-            print(f"Saved conversation stats for {conv_id} to {conv_stats_path}")
+            raise ValueError(f"Unsupported update method: {method}")
 
-    def _write_user_search_results(self, user_search_path, search_results, conv_id):
-        """Write per-user search results to a temporary JSON file."""
-        with open(user_search_path, "w") as fw:
-            json.dump(dict(search_results), fw, indent=2)
-            print(f"Save search results {conv_id}")
+    def _update_context_template(self, conv_id, query, answer):
+        new_context = f"User: {query}\nAssistant: {answer}\n\n"
+        if self.pre_context_cache[conv_id] is None:
+            self.pre_context_cache[conv_id] = ""
+        self.pre_context_cache[conv_id] += new_context
 
     def _process_single_qa(
         self,
@@ -136,24 +88,35 @@ def _process_single_qa(
         conv_stats,
     ):
         query = qa.get("question")
+        gold_answer = qa.get("answer")
         qa_category = qa.get("category")
         if qa_category == 5:
             return None
 
         # Search
-        context, search_duration_ms = self.search_query(
+        cur_context, search_duration_ms = self.search_query(
             client, query, metadata, frame, reversed_client=reversed_client, top_k=top_k
         )
-        if not context:
+        if not cur_context:
             logger.warning(f"No context found for query: {query[:100]}")
-            context = ""
+            cur_context = ""
 
         # Context answerability analysis (for memos_scheduler only)
-        gold_answer = qa.get("answer")
         if self.pre_context_cache[conv_id] is None:
             # Update pre-context cache with current context
-            with self.stats_lock:
-                self.pre_context_cache[conv_id] = context
+            if self.frame in [MEMOS_MODEL, MEMOS_SCHEDULER_MODEL]:
+                self.update_context(
+                    conv_id=conv_id,
+                    method=self.context_update_method,
+                    cur_context=cur_context,
+                )
+            else:
+                self.update_context(
+                    conv_id=conv_id,
+                    method=self.context_update_method,
+                    query=query,
+                    answer=gold_answer,
+                )
             return None
 
         can_answer = False
@@ -181,15 +144,9 @@ def _process_single_qa(
             )
             self.save_stats()
 
-        # Update pre-context cache with current context
-        with self.stats_lock:
-            self.pre_context_cache[conv_id] = context
-
-        self.print_eval_info()
-
         # Generate answer
         answer_start = time()
-        answer = self.locomo_response(frame, oai_client, context, query)
+        answer = self.locomo_response(frame, oai_client, self.pre_context_cache[conv_id], query)
         response_duration_ms = (time() - answer_start) * 1000
 
         # Record case for memos_scheduler
@@ -199,7 +156,7 @@ def _process_single_qa(
                     conv_id=conv_id,
                     query=query,
                     answer=answer,
-                    context=context,
+                    context=cur_context,
                     pre_context=self.pre_context_cache[conv_id],
                     can_answer=can_answer,
                     can_answer_reason=f"Context analysis result: {'can answer' if can_answer else 'cannot answer'}",
@@ -248,146 +205,37 @@ def _process_single_qa(
 
         logger.info(f"Processed question: {query[:100]}")
         logger.info(f"Answer: {answer[:100]}")
+
+        # Update pre-context cache with current context
+        with self.stats_lock:
+            if self.frame in [MEMOS_MODEL, MEMOS_SCHEDULER_MODEL]:
+                self.update_context(
+                    conv_id=conv_id,
+                    method=self.context_update_method,
+                    cur_context=cur_context,
+                )
+            else:
+                self.update_context(
+                    conv_id=conv_id,
+                    method=self.context_update_method,
+                    query=query,
+                    answer=gold_answer,
+                )
+
+        self.print_eval_info()
+
         return {
             "question": query,
             "answer": answer,
             "category": qa_category,
             "golden_answer": gold_answer,
-            "search_context": context,
+            "search_context": cur_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,
         }
 
-    def process_user(self, conv_id, locomo_df, frame, version, top_k=20):
-        user_search_path = self.result_dir / f"tmp/{frame}_locomo_search_results_{conv_id}.json"
-        user_search_path.parent.mkdir(exist_ok=True, parents=True)
-        search_results = defaultdict(list)
-        response_results = defaultdict(list)
-        conv_stats_path = self.stats_dir / f"{frame}_{version}_conv_{conv_id}_stats.json"
-
-        conversation = locomo_df["conversation"].iloc[conv_id]
-        speaker_a = conversation.get("speaker_a", "speaker_a")
-        speaker_b = conversation.get("speaker_b", "speaker_b")
-
-        # Use temporal_locomo data if available, otherwise fall back to original locomo data
-        temporal_conv = self.temporal_locomo_data[conv_id]
-        conv_id = temporal_conv["conversation_id"]
-        speaker_a_user_id = f"{conv_id}_speaker_a"
-        speaker_b_user_id = f"{conv_id}_speaker_b"
-
-        # Process temporal data by days
-        day_groups = {}
-        for day_id, day_data in temporal_conv["days"].items():
-            day_groups[day_id] = day_data["qa_pairs"]
-
-        # Initialize conversation-level statistics
-        conv_stats = self._initialize_conv_stats()
-
-        metadata = self._build_metadata(
-            speaker_a, speaker_b, speaker_a_user_id, speaker_b_user_id, conv_id
-        )
-
-        client, reversed_client = self._get_clients(
-            frame, speaker_a_user_id, speaker_b_user_id, conv_id, version, top_k
-        )
-
-        oai_client = OpenAI(api_key=self.openai_api_key, base_url=self.openai_base_url)
-
-        with self.stats_lock:
-            self.pre_context_cache[conv_id] = None
-
-        def process_qa(qa):
-            return self._process_single_qa(
-                qa,
-                client=client,
-                reversed_client=reversed_client,
-                metadata=metadata,
-                frame=frame,
-                version=version,
-                conv_id=conv_id,
-                conv_stats_path=conv_stats_path,
-                oai_client=oai_client,
-                top_k=top_k,
-                conv_stats=conv_stats,
-            )
-
-        # ===================================
-        conv_stats["theoretical_total_queries"] = 0
-        for day, qa_list in day_groups.items():
-            conv_stats["theoretical_total_queries"] += len(qa_list) - 1
-            conv_stats["processing_failure_count"] = 0
-            print(f"Processing user {conv_id} day {day}")
-            for qa in tqdm(qa_list, desc=f"Processing user {conv_id} day {day}"):
-                try:
-                    result = process_qa(qa)
-                except Exception as e:
-                    logger.error(f"Error: {e}. traceback: {traceback.format_exc()}")
-                    conv_stats["processing_failure_count"] += 1
-                    continue
-                if result:
-                    context_preview = (
-                        result["search_context"][:20] + "..."
-                        if result["search_context"]
-                        else "No context"
-                    )
-                    if "can_answer" in result:
-                        logger.info("Print can_answer case")
-                        logger.info(
-                            {
-                                "question": result["question"][:100],
-                                "pre context can answer": result["can_answer"],
-                                "answer": result["answer"][:100],
-                                "golden_answer": result["golden_answer"],
-                                "search_context": context_preview[:100],
-                                "search_duration_ms": result["search_duration_ms"],
-                            }
-                        )
-
-                    search_results[conv_id].append(
-                        {
-                            "question": result["question"],
-                            "context": result["search_context"],
-                            "search_duration_ms": result["search_duration_ms"],
-                        }
-                    )
-                    response_results[conv_id].append(result)
-
-            logger.warning(
-                f"Finished processing user {conv_id} day {day}, data_length: {len(qa_list)}"
-            )
-
-        # recording separate search results
-        with open(user_search_path, "w") as fw:
-            json.dump(dict(search_results), fw, indent=2)
-            print(f"Save search results {conv_id}")
-
-        # Dump stats after processing each user
-        self.save_stats()
-
-        return search_results, response_results
-
-    def process_user_wrapper(self, args):
-        """
-        Wraps the process_user function to support parallel execution and error handling.
-
-        Args:
-            args: Tuple containing parameters for process_user
-
-        Returns:
-            tuple: Contains user results or error information
-        """
-        idx, locomo_df, frame, version, top_k = args
-        try:
-            print(f"Processing user {idx}...")
-            user_search_results, user_response_results = self.process_user(
-                idx, locomo_df, frame, version, top_k
-            )
-            return (user_search_results, user_response_results, None)
-        except Exception as e:
-            return (None, None, (idx, e, traceback.format_exc()))
-
     def run_locomo_processing(self, num_users=10):
         load_dotenv()