Fix HF warnings (attention_mask, pad_token_id) and add DEBUG mode for fast eval.

files changed: eval.py, metrics/knowmem.py, metrics/privleak.py, metrics/verbmem.py

Author: liranc6

MUSE/eval.py

@@ -1,4 +1,4 @@
-from metrics.verbmem import eval as eval_verbmem
+from metrics.verbmem import eval as eval_ve rbmem
 from metrics.privleak import eval as eval_privleak
 from metrics.knowmem import eval as eval_knowmem
 from utils import load_model, load_tokenizer, write_csv, read_json, write_json
@@ -29,6 +29,7 @@ def eval_model(
     knowmem_retain_qa_file: str | None = None,
     knowmem_retain_qa_icl_file: str | None = None,
     temp_dir: str | None = None,
+    DEBUG: bool = False,
 ) -> Dict[str, float]:
     # Argument sanity check
     if not metrics:
@@ -50,10 +51,13 @@ def eval_model(
 
     out = {}
     model = model.to('cuda')
+    debug_subset_len = 3 if DEBUG else None
 
     # 1. verbmem_f
     if 'verbmem_f' in metrics:
         data = read_json(verbmem_forget_file)
+        if DEBUG:
+            data = data[:debug_subset_len]
         agg, log = eval_verbmem(
             prompts=[d['prompt'] for d in data],
             gts=[d['gt'] for d in data],
@@ -67,10 +71,17 @@ def eval_model(
 
     # 2. privleak
     if 'privleak' in metrics:
+        forget_data = read_json(privleak_forget_file)
+        retain_data = read_json(privleak_retain_file)
+        holdout_data = read_json(privleak_holdout_file)
+        if DEBUG:
+            forget_data = forget_data[:debug_subset_len]
+            retain_data = retain_data[:debug_subset_len]
+            holdout_data = holdout_data[:debug_subset_len]
         auc, log = eval_privleak(
-            forget_data=read_json(privleak_forget_file),
-            retain_data=read_json(privleak_retain_file),
-            holdout_data=read_json(privleak_holdout_file),
+            forget_data=forget_data,
+            retain_data=retain_data,
+            holdout_data=holdout_data,
             model=model, tokenizer=tokenizer
         )
         if temp_dir is not None:
@@ -82,6 +93,9 @@ def eval_model(
     if 'knowmem_f' in metrics:
         qa = read_json(knowmem_forget_qa_file)
         icl = read_json(knowmem_forget_qa_icl_file)
+        if DEBUG:
+            qa = qa[:debug_subset_len]
+            icl = icl[:debug_subset_len]
         agg, log = eval_knowmem(
             questions=[d['question'] for d in qa],
             answers=[d['answer'] for d in qa],
@@ -99,6 +113,9 @@ def eval_model(
     if 'knowmem_r' in metrics:
         qa = read_json(knowmem_retain_qa_file)
         icl = read_json(knowmem_retain_qa_icl_file)
+        if DEBUG:
+            qa = qa[:debug_subset_len]
+            icl = icl[:debug_subset_len]
         agg, log = eval_knowmem(
             questions=[d['question'] for d in qa],
             answers=[d['answer'] for d in qa],
@@ -122,7 +139,8 @@ def load_then_eval_models(
     tokenizer_dir: str = LLAMA_DIR,
     out_file: str | None = None,
     metrics: List[str] = SUPPORTED_METRICS,
-    temp_dir: str = "temp"
+    temp_dir: str = "temp",
+    DEBUG: bool = False,
 ) -> DataFrame:
     print(out_file)
     # Argument sanity check
@@ -140,7 +158,8 @@ def load_then_eval_models(
         tokenizer = load_tokenizer(tokenizer_dir)
         res = eval_model(
             model, tokenizer, metrics, corpus,
-            temp_dir=os.path.join(temp_dir, name)
+            temp_dir=os.path.join(temp_dir, name),
+            DEBUG=DEBUG
         )
         out.append({'name': name} | res)
         if out_file is not None: write_csv(out, out_file)

MUSE/metrics/knowmem.py

@@ -28,18 +28,23 @@ def eval(
     for question, answer in tzip(questions, answers):
         prompt = general_prompt + f"Question: {question}\nAnswer: "
 
-        # Encode the `prompt` into `input_ids`
-        input_ids = tokenizer(
+        # Encode the `prompt` into `input_ids` and `attention_mask`
+        inputs = tokenizer(
             prompt,
             return_tensors='pt',
-            add_special_tokens=True).input_ids
+            add_special_tokens=True
+        )
+        input_ids = inputs.input_ids
+        attention_mask = inputs.attention_mask
 
         # Use the `model` to generate the continuation of the `input_ids`.
         output_ids = model.generate(
-            input_ids.to(model.device),
+            input_ids=input_ids.to(model.device),
+            attention_mask=attention_mask.to(model.device),
             max_new_tokens=max_new_tokens,
             do_sample=False,
-            pad_token_id=tokenizer.pad_token_id)
+            pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
+        )
         output_ids = output_ids[:, len(input_ids[0]):]
 
         output = tokenizer.batch_decode(

MUSE/metrics/privleak.py

@@ -11,15 +11,21 @@
 
 
 def compute_ppl(text: str, model, tokenizer, device='cuda'):
-    input_ids = torch.tensor(tokenizer.encode(text)).unsqueeze(0)
-    input_ids = input_ids.to(device)
+    # Tokenize with attention_mask and padding
+    inputs = tokenizer(
+        text,
+        return_tensors='pt',
+        add_special_tokens=True
+    )
+    input_ids = inputs['input_ids'].to(device)
+    attention_mask = inputs['attention_mask'].to(device)
     with torch.no_grad():
-        outputs = model(input_ids, labels=input_ids)
+        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
     loss, logits = outputs[:2]
 
     probabilities = torch.nn.functional.log_softmax(logits, dim=-1)
-    all_prob = []
     input_ids_processed = input_ids[0][1:]
+    all_prob = []
     for i, token_id in enumerate(input_ids_processed):
         probability = probabilities[0, i, token_id].item()
         all_prob.append(probability)

MUSE/metrics/verbmem.py

@@ -11,22 +11,26 @@ def eval(
 ):
     logger = RougeEvalLogger()
     for prompt, gt in tzip(prompts, gts):
-        # Encode the `prompt` into `input_ids`
-        input_ids = tokenizer(
+        # Encode the `prompt` into `input_ids` and `attention_mask`
+        inputs = tokenizer(
             prompt,
             return_tensors='pt',
             add_special_tokens=True
-        ).input_ids
+        )
+        input_ids = inputs.input_ids
+        attention_mask = inputs.attention_mask
         assert len(input_ids) == 1
 
         gt_ids = tokenizer(gt, return_tensors='pt', add_special_tokens=True).input_ids[:, :max_new_tokens]
 
         # Use the `model` to generate the continuation of the `input_ids`.
         output_ids = model.generate(
-            input_ids.to(model.device),
+            input_ids=input_ids.to(model.device),
+            attention_mask=attention_mask.to(model.device),
             max_new_tokens=max_new_tokens,
             do_sample=False,
-            pad_token_id=tokenizer.pad_token_id)
+            pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
+        )
         output_ids = output_ids[:, len(input_ids[0]):]
         output = tokenizer.batch_decode(
             output_ids,