Fix generated evaluation consistency and preserve locality evidence

easyeditor/editors/concept_editor.py

@@ -237,8 +237,6 @@ def edit(self,
                         f"locality.{locality_key}",
                         build_locality_metric_meta(locality_key, self.hparams, self.model_name),
                     )
-                    all_metrics[i]['post']['locality'].pop(f'{locality_key}_output')
-                all_metrics[i]['pre'].pop('locality')
 
             LOG.info(f"Evaluation took {time() - start}")
 

easyeditor/editors/editor.py

@@ -51,6 +51,49 @@ def seed_everything(seed):
     random.seed(seed)
 
 seed_everything(42)
+
+
+def finalize_locality_metrics(metric, request, hparams, model_name):
+    if not request.get("locality"):
+        return
+
+    pre_locality = metric["pre"]["locality"]
+    post_locality = metric["post"]["locality"]
+    evaluation_type = getattr(hparams, "evaluation_type", None)
+    uses_generated_text = evaluation_type in ["LLM-judge", "generate-text"]
+
+    for locality_key in request["locality"]:
+        if uses_generated_text:
+            output_key = f"{locality_key}_gen_content"
+        else:
+            output_key = f"{locality_key}_output"
+
+        pre_outputs = pre_locality[output_key]
+        post_outputs = post_locality[output_key]
+        if len(pre_outputs) != len(post_outputs):
+            raise ValueError(
+                f"Locality output count mismatch for `{locality_key}`: "
+                f"{len(pre_outputs)} pre-edit outputs and "
+                f"{len(post_outputs)} post-edit outputs."
+            )
+
+        if uses_generated_text:
+            locality_result = [
+                float(pre_output == post_output)
+                for pre_output, post_output in zip(pre_outputs, post_outputs)
+            ]
+        else:
+            locality_result = [
+                float(np.mean(np.equal(pre_output, post_output)))
+                for pre_output, post_output in zip(pre_outputs, post_outputs)
+            ]
+
+        post_locality[f"{locality_key}_acc"] = locality_result
+        attach_metric_meta(
+            metric["post"],
+            f"locality.{locality_key}",
+            build_locality_metric_meta(locality_key, hparams, model_name),
+        )
 
 class BaseEditor:
     """Base editor for all methods"""
@@ -252,22 +295,12 @@ def batch_edit(self,
                 restore_after_edit(self, edited_model, weights_copy)
 
             for i, request in enumerate(record_chunks):
-                if 'locality' in chunk_metrics[i]['post'].keys():
-                    for locality_key in request['locality'].keys():
-                        locality_result = []
-                        if hasattr(self.hparams, 'evaluation_type') and self.hparams.evaluation_type == "LLM-judge":
-                            locality_result.append(float(chunk_metrics[i]['post']['locality'][f'{locality_key}_output']==chunk_metrics[i]['pre']['locality'][f'{locality_key}_output']))
-                        else:
-                            for ans, label in zip(chunk_metrics[i]['post']['locality'][f'{locality_key}_output'], chunk_metrics[i]['pre']['locality'][f'{locality_key}_output']):
-                                locality_result.append(np.mean(np.equal(ans, label)))
-                        chunk_metrics[i]['post']['locality'][f'{locality_key}_acc'] = locality_result
-                        attach_metric_meta(
-                            chunk_metrics[i]['post'],
-                            f"locality.{locality_key}",
-                            build_locality_metric_meta(locality_key, self.hparams, self.model_name),
-                        )
-                        chunk_metrics[i]['post']['locality'].pop(f'{locality_key}_output')
-                    chunk_metrics[i]['pre'].pop('locality')
+                finalize_locality_metrics(
+                    chunk_metrics[i],
+                    request,
+                    self.hparams,
+                    self.model_name,
+                )
 
                 if verbose:
                     LOG.info(
@@ -356,22 +389,12 @@ def edit_evaluation(all_metrics, request, edited_model, idx, test_generation, ic
                 })
                 if "metric_kwargs" in kwargs:
                     all_metrics[idx].update(compute_sent_metric(self.model, edited_model, self.model_name, self.hparams, self.tok,metric_kwargs=kwargs["metric_kwargs"][idx], device=self.hparams.device))
-                if 'locality' in all_metrics[idx]['post'].keys() and not hasattr(self.hparams, 'evaluation_type'):
-                    for locality_key in request['locality'].keys():
-                        locality_result = []
-                        if hasattr(self.hparams, 'evaluation_type'):
-                            locality_result.append(float(all_metrics[idx]['post']['locality'][f'{locality_key}_output']==all_metrics[idx]['pre']['locality'][f'{locality_key}_output']))
-                        else:
-                            for ans, label in zip(all_metrics[idx]['post']['locality'][f'{locality_key}_output'], all_metrics[idx]['pre']['locality'][f'{locality_key}_output']):
-                                locality_result.append(np.mean(np.equal(ans, label)))
-                        all_metrics[idx]['post']['locality'][f'{locality_key}_acc'] = locality_result
-                        attach_metric_meta(
-                            all_metrics[idx]['post'],
-                            f"locality.{locality_key}",
-                            build_locality_metric_meta(locality_key, self.hparams, self.model_name),
-                        )
-                        all_metrics[idx]['post']['locality'].pop(f'{locality_key}_output')
-                    all_metrics[idx]['pre'].pop('locality')
+                finalize_locality_metrics(
+                    all_metrics[idx],
+                    request,
+                    self.hparams,
+                    self.model_name,
+                )
 
             if verbose:
                 LOG.info(f"{idx} editing: {request['prompt']} -> {request['target_new']}  \n\n {all_metrics[idx]}")
@@ -395,8 +418,7 @@ def edit_evaluation(all_metrics, request, edited_model, idx, test_generation, ic
 
         if isinstance(edited_model, LORA):
             edited_model = edited_model.model
-        if not hasattr(self.hparams, 'evaluation_type') or self.hparams.evaluation_type != "generate-text":
-            summary_metrics(all_metrics)
+        summary_metrics(all_metrics)
 
         return all_metrics, edited_model, weights_copy
 

easyeditor/editors/multimodal_editor.py

@@ -53,10 +53,42 @@ def make_logs():
     LOG.addHandler(s_h)
 
 
+def get_aligned_topk_token_ids(pre_logits, post_logits, k):
+    pre_logits = torch.as_tensor(pre_logits)
+    post_logits = torch.as_tensor(post_logits)
+    if pre_logits.dim() != 3 or post_logits.dim() != 3:
+        raise ValueError(
+            "pre_logits and post_logits must be rank-3 tensors shaped "
+            "[batch, sequence, vocab]."
+        )
+    if pre_logits.shape[0] != post_logits.shape[0]:
+        raise ValueError(
+            "pre_logits and post_logits must have the same batch size: "
+            f"{pre_logits.shape[0]} != {post_logits.shape[0]}"
+        )
+
+    if post_logits.shape[1] > pre_logits.shape[1]:
+        post_logits = post_logits[:, -pre_logits.shape[1]:, :]
+    else:
+        pre_logits = pre_logits[:, -post_logits.shape[1]:, :]
+
+    k = min(k, pre_logits.shape[-1], post_logits.shape[-1])
+    pre_topk = torch.topk(pre_logits, k=k, dim=-1).indices.cpu().tolist()
+    post_topk = torch.topk(post_logits, k=k, dim=-1).indices.cpu().tolist()
+    return pre_topk, post_topk
+
+
 def attach_topk_locality_metrics(metrics):
     if 'locality_output' in metrics['post'].keys():
         assert len(metrics['post']['locality_output']) == \
                 len(metrics['pre']['locality_output'])
+        pre_topk, post_topk = get_aligned_topk_token_ids(
+            metrics['pre']['locality_output'],
+            metrics['post']['locality_output'],
+            k=1,
+        )
+        metrics['pre']['locality_topk_tokens'] = pre_topk
+        metrics['post']['locality_topk_tokens'] = post_topk
         metrics['post']['locality_acc'] = topk_token_match(
             metrics['pre']['locality_output'],
             metrics['post']['locality_output'],
@@ -78,6 +110,13 @@ def attach_topk_locality_metrics(metrics):
     if 'multimodal_locality_output' in metrics['post'].keys():
         assert len(metrics['post']['multimodal_locality_output']) == \
                 len(metrics['pre']['multimodal_locality_output'])
+        pre_topk, post_topk = get_aligned_topk_token_ids(
+            metrics['pre']['multimodal_locality_output'],
+            metrics['post']['multimodal_locality_output'],
+            k=10,
+        )
+        metrics['pre']['multimodal_locality_topk_tokens'] = pre_topk
+        metrics['post']['multimodal_locality_topk_tokens'] = post_topk
         metrics['post']['multimodal_locality_acc'] = topk_token_match(
             metrics['pre']['multimodal_locality_output'],
             metrics['post']['multimodal_locality_output'],

easyeditor/evaluate/evaluate.py

@@ -16,6 +16,7 @@
 from ..util import HyperParams
 from ..util.device import normalize_device
 from .evaluate_utils import (
+    generate_texts,
     test_seq2seq_batch_prediction_acc, 
     test_batch_prediction_acc, 
     test_prediction_acc,
@@ -126,7 +127,7 @@ def compute_rewrite_or_rephrase_quality(
             f"{key}_gen_content": gen_content
         }
     elif hasattr(hparams, 'evaluation_type') and hparams.evaluation_type == "generate-text":
-        gen_content_model = test_prediction_acc_LLM_judge(model, tok, hparams, prompt, target_new, device, locality=False)
+        gen_content_model = generate_texts(model, tok, hparams, prompt, device)
         ret = {
             f"{key}_gen_content": gen_content_model
         }
@@ -176,8 +177,13 @@ def compute_locality_quality(
 ) -> typing.Dict:
 
     # using real-world evaluation: autoregressive decoding, natural stop criteria, LLM-as-a-Judge
-    if hasattr(hparams, 'evaluation_type'):
-        loc_tokens = test_prediction_acc_LLM_judge(model, tok, hparams, prompt, locality_ground_truth, device, locality=True)
+    evaluation_type = getattr(hparams, "evaluation_type", None)
+    if evaluation_type in ["LLM-judge", "generate-text"]:
+        return {
+            f"{locality_key}_gen_content": generate_texts(
+                model, tok, hparams, prompt, device
+            )
+        }
     else:  # traditional evaluation 
         if 't5' in model_name.lower():
             loc_tokens = test_seq2seq_batch_prediction_acc(model, tok, hparams, prompt, locality_ground_truth, device, locality=True)
@@ -202,10 +208,21 @@ def compute_portability_quality(
     device,
 ) -> typing.Dict:
     # using real-world evaluation: autoregressive decoding, natural stop criteria, LLM-as-a-Judge
-    if hasattr(hparams, 'evaluation_type') and hparams.evaluation_type == "LLM-judge":
-        portability_correct = test_prediction_acc_LLM_judge(model, tok, hparams, prompt, ground_truth, device, locality=False)
-    elif hasattr(hparams, 'evaluation_type') and hparams.evaluation_type == "generate-text":
-        portability_correct = test_prediction_acc_LLM_judge(model, tok, hparams, prompt, ground_truth, device, locality=False)
+    evaluation_type = getattr(hparams, "evaluation_type", None)
+    if evaluation_type == "LLM-judge":
+        portability_correct, gen_content = test_prediction_acc_LLM_judge(
+            model, tok, hparams, prompt, ground_truth, device, locality=False
+        )
+        return {
+            f"{portability_key}_acc": portability_correct,
+            f"{portability_key}_gen_content": gen_content,
+        }
+    elif evaluation_type == "generate-text":
+        return {
+            f"{portability_key}_gen_content": generate_texts(
+                model, tok, hparams, prompt, device
+            )
+        }
     else:  # traditional evaluation
         if 't5' in model_name.lower():
             portability_correct = test_seq2seq_batch_prediction_acc(model, tok, hparams, prompt, ground_truth, device)

easyeditor/evaluate/evaluate_utils.py

@@ -101,12 +101,11 @@ def llm_judge(question, ground_truth, prediction, api_key):
     time.sleep(1) # avoid high rate of request
     return llm_score
 
-def test_prediction_acc_LLM_judge(model, tok, hparams, prompts, targets, device, locality=False):
-    # generation & truncation
-    all_score = []
-    all_response = []
+def generate_texts(model, tok, hparams, prompts, device):
     if isinstance(prompts, str):
-        prompts, targets = [prompts, ], [targets, ]
+        prompts = [prompts]
+
+    all_response = []
     for prompt in prompts:
         messages = [
                 {"role": "system", "content": "You are a helpful assistant."},
@@ -146,25 +145,32 @@ def test_prediction_acc_LLM_judge(model, tok, hparams, prompts, targets, device,
         gen_content = tok.decode(trunc_gen_tokens)
         suffixes_to_remove = [".", "\n", tok.eos_token]
         for suffix in suffixes_to_remove:
-            if gen_content.endswith(suffix):
+            if suffix and gen_content.endswith(suffix):
                 gen_content = gen_content[:-len(suffix)]
-        # LLM-as-a-Judge
-        if hparams.evaluation_type == "generate-text":
-            all_response.append(gen_content)
-        elif hparams.evaluation_type == "LLM-judge" and hasattr(hparams, 'api_key') and hparams.api_key:
-            LLM_Score = llm_judge(prompts, targets, gen_content, hparams.api_key)
-            all_score.append(LLM_Score)
-            all_response.append(gen_content)
+
+        all_response.append(gen_content)
+
+    return all_response
+
+
+def test_prediction_acc_LLM_judge(model, tok, hparams, prompts, targets, device, locality=False):
+    if isinstance(prompts, str):
+        prompts = [prompts]
+    if isinstance(targets, str):
+        targets = [targets]
+    if len(prompts) != len(targets):
+        raise ValueError("The number of prompts and targets must match.")
+
+    all_response = generate_texts(model, tok, hparams, prompts, device)
+    all_score = []
+    for prompt, target, gen_content in zip(prompts, targets, all_response):
+        if hasattr(hparams, 'api_key') and hparams.api_key:
+            score = llm_judge(prompt, target, gen_content, hparams.api_key)
         else:
-            # the user do not provide api key, using exact match as an alternative
-            EM_Score = float(exact_match_score(gen_content, targets))
-            all_score.append(EM_Score)
-            all_response.append(gen_content)
-
-    if len(all_score) > 0:
-        return all_score, all_response
-    else:
-        return all_response
+            score = float(exact_match_score(gen_content, target))
+        all_score.append(score)
+
+    return all_score, all_response
 
 def test_batch_prediction_acc(model, tok, hparams, prompts, target, device, locality=False):
     prompt_tok = tok(