refactor: rewrite IFD evaluation with clean three-layer architecture

scripts/eval/evaluate_ifd.py

@@ -1,31 +1,14 @@
 """IFD (Instruction Following Difficulty) data quality scoring.
 
-Computes IFD scores for instruction-response pairs to guide data selection.
+IFD = conditional_NLL / unconditional_NLL
-IFD = conditional_NLL / unconditional_NLL, where:
 
-- conditional_NLL: average CE loss on response tokens given instruction context
+- Messages format: plain text concatenation (no chat template)
-- unconditional_NLL: average CE loss on response tokens alone
+- Plain format: raw instr_key + resp_key fields
-
-Higher IFD (close to 1) = instruction provides less help = harder sample.
-Lower IFD (close to 0) = instruction provides strong guidance = easy sample.
-IFD > 1 = instruction misleads the model = likely low-quality data.
-
-Usage::
-
-    python scripts/eval/ifd.py --param_path ./params \
-        --input data.jsonl --output data_with_ifd.jsonl \
-        --instr_key instruction --resp_key response
-
-Disable chat template::
-
-    python scripts/eval/ifd.py --param_path ./params \
-        --input data.jsonl --output data_with_ifd.jsonl \
-        --instr_key instruction --resp_key response \
-        --no_chat_template
 """
 
 import argparse
 import json
+import statistics
 
 import torch
 import torch.nn.functional as F
@@ -35,172 +18,88 @@ from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 
 
-def compute_ifd(
+def _score(context_ids, resp_ids, model, device):
-    model,
+    """Core IFD computation: context → L_cond, response alone → L_uncond."""
-    tokenizer,
-    instruction: str,
-    response: str,
-    device: str,
-    max_len: int = 2048,
-    use_chat_template: bool = False,
-) -> dict:
-    if use_chat_template:
-        return _compute_ifd_with_template(
-            model, tokenizer, instruction, response, device, max_len
-        )
-    return _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len)
-
-
-def _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len) -> dict:
-    instr_ids = tokenizer.encode(instruction, add_special_tokens=False)
-    resp_ids = tokenizer.encode(response, add_special_tokens=False)
-
-    if len(resp_ids) > max_len:
-        resp_ids = resp_ids[:max_len]
-
     if not resp_ids:
-        return {
+        return None
-            "L_cond": None,
+    full_ids = context_ids + resp_ids
-            "L_uncond": None,
+    inp_full = torch.tensor([full_ids], device=device, dtype=torch.long)
-            "ifd": None,
+    inp_resp = torch.tensor([resp_ids], device=device, dtype=torch.long)
-            "error": "empty response",
+    logits_full = model(inp_full)["logits"][0]
-        }
+    logits_resp = model(inp_resp)["logits"][0]
-
+    ctx_len = len(context_ids)
-    qa_len = len(instr_ids) + len(resp_ids)
+    resp_logits = logits_full[ctx_len - 1 : -1]
-    if qa_len > max_len:
+    resp_targets = torch.tensor(resp_ids, device=device, dtype=torch.long)
-        overflow = qa_len - max_len
-        if overflow >= len(instr_ids):
-            resp_ids = resp_ids[:max_len]
-            instr_ids = []
-        else:
-            instr_ids = instr_ids[overflow:]
-
-    if not instr_ids:
-        return {
-            "L_cond": None,
-            "L_uncond": None,
-            "ifd": None,
-            "error": "response too long for context",
-        }
-
-    instr_len = len(instr_ids)
-    resp_len = len(resp_ids)
-
-    qa_ids = instr_ids + resp_ids
-
-    with torch.inference_mode():
-        logits_qa = model(torch.tensor([qa_ids], device=device, dtype=torch.long))[
-            "logits"
-        ][0]
-        logits_resp = model(torch.tensor([resp_ids], device=device, dtype=torch.long))[
-            "logits"
-        ][0]
-
-    resp_logits = logits_qa[instr_len - 1 : -1]
-    resp_targets = logits_resp.new_tensor(resp_ids, dtype=torch.long)
     L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()
-
     unp_logits = logits_resp[:-1]
-    unp_targets = logits_resp.new_tensor(resp_ids[1:], dtype=torch.long)
+    unp_targets = torch.tensor(resp_ids[1:], device=device, dtype=torch.long)
     L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
-
     ifd = L_cond / L_uncond if L_uncond > 0 else None
-
     return {
         "L_cond": round(L_cond, 6),
         "L_uncond": round(L_uncond, 6),
         "ifd": round(ifd, 6) if ifd is not None else None,
-        "instr_len": instr_len,
-        "resp_len": resp_len,
-        "error": None,
-    }
-
-
-def _compute_ifd_with_template(
-    model, tokenizer, instruction, response, device, max_len
-) -> dict:
-    instr_prefix = tokenizer.apply_chat_template(
-        [{"role": "user", "content": instruction}],
-        tokenize=False,
-        add_generation_prompt=True,
-    )
-    full_text = tokenizer.apply_chat_template(
-        [
-            {"role": "user", "content": instruction},
-            {"role": "assistant", "content": response},
-        ],
-        tokenize=False,
-        add_generation_prompt=False,
-    )
-
-    full_ids = tokenizer.encode(full_text)
-    prefix_ids = tokenizer.encode(instr_prefix)
-    resp_ids = tokenizer.encode(response)
-
-    if not resp_ids:
-        return {
-            "L_cond": None,
-            "L_uncond": None,
-            "ifd": None,
-            "error": "empty response",
-        }
-
-    if len(full_ids) > max_len:
-        overflow = len(full_ids) - max_len
-        full_ids = full_ids[overflow:]
-        prefix_len = len(prefix_ids) - overflow
-        prefix_len = max(0, prefix_len)
-    else:
-        prefix_len = len(prefix_ids)
-
-    cond_tensor = torch.tensor([full_ids], device=device, dtype=torch.long)
-
-    with torch.inference_mode():
-        logits_qa = model(cond_tensor)["logits"][0]
-
-    resp_start = prefix_len - 1
-    resp_end = len(full_ids) - 1
-    if resp_end <= resp_start:
-        return {
-            "L_cond": None,
-            "L_uncond": None,
-            "ifd": None,
-            "error": "response truncated entirely",
-        }
-
-    resp_logits = logits_qa[resp_start:resp_end]
-    resp_targets = torch.tensor(full_ids[prefix_len:], device=device, dtype=torch.long)
-    L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()
-
-    resp_tensor = torch.tensor([resp_ids], device=device, dtype=torch.long)
-
-    with torch.inference_mode():
-        logits_resp = model(resp_tensor)["logits"][0]
-
-    unp_logits = logits_resp[:-1]
-    unp_targets = resp_tensor[0, 1:]
-    L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
-
-    ifd = L_cond / L_uncond if L_uncond > 0 else None
-
-    return {
-        "L_cond": round(L_cond, 6),
-        "L_uncond": round(L_uncond, 6),
-        "ifd": round(ifd, 6) if ifd is not None else None,
-        "instr_len": prefix_len,
         "resp_len": len(resp_ids),
-        "error": None,
     }
 
 
+def _trim(context_ids, resp_ids, max_len):
+    """Truncate to fit max_len, keeping response intact if possible."""
+    if len(resp_ids) > max_len // 2:
+        resp_ids = resp_ids[: max_len // 2]
+    full_ids = context_ids + resp_ids
+    if len(full_ids) <= max_len:
+        return context_ids, resp_ids
+    overflow = len(full_ids) - max_len
+    if overflow >= len(context_ids):
+        return [], resp_ids[:max_len]
+    return context_ids[overflow:], resp_ids
+
+
+def score_plain(model, tokenizer, instruction, response, device, max_len=2048):
+    """Compute IFD for a single instruction-response pair (plain format)."""
+    ctx_ids = tokenizer.encode(instruction, add_special_tokens=False)
+    resp_ids = tokenizer.encode(response, add_special_tokens=False)
+    ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
+    if not ctx_ids or not resp_ids:
+        return {"L_cond": None, "L_uncond": None, "ifd": None, "error": "empty"}
+    return _score(ctx_ids, resp_ids, model, device)
+
+
+def score_messages(model, tokenizer, messages, device, max_len=2048):
+    """Compute IFD for each assistant turn in a messages array."""
+    turns = []
+    for i, msg in enumerate(messages):
+        if msg.get("role") != "assistant":
+            continue
+        ctx_text = "\n\n".join(m["content"] for m in messages[:i])
+        ctx_ids = tokenizer.encode(ctx_text)
+        resp_ids = tokenizer.encode(msg["content"], add_special_tokens=False)
+        ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
+        if ctx_ids and resp_ids:
+            turns.append(_score(ctx_ids, resp_ids, model, device))
+    if not turns:
+        return None
+    valid = [t for t in turns if t is not None and t["ifd"] is not None]
+    if not valid:
+        return {"ifd": None, "ifd_turns": turns}
+    avg = sum(t["ifd"] for t in valid) / len(valid)
+    return {
+        "ifd": avg,
+        "ifd_detail": valid[0] if len(valid) == 1 else None,
+        "ifd_turns": turns,
+    }
+
+
+@torch.inference_mode()
 def process_file(
-    param_path: str,
+    param_path,
-    input_file: str,
+    input_file,
-    output_file: str,
+    output_file,
-    instr_key: str,
+    instr_key,
-    resp_key: str,
+    resp_key,
-    max_len: int = 2048,
+    max_len=2048,
-    use_chat_template: bool = False,
+    data_format="plain",
 ):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.bfloat16 if device == "cuda" else torch.float32
@@ -210,42 +109,35 @@ def process_file(
     model.to(device=device, dtype=dtype)
     model.eval()
 
-    if use_chat_template and tokenizer._chat_template is None:
+    with open(input_file, encoding="utf-8") as f:
-        raise RuntimeError(
-            "--use_chat_template specified but tokenizer has no chat template. "
-            "Add a chat_template to tokenizer_config.json or omit the flag."
-        )
-
-    with open(input_file, "r", encoding="utf-8") as f:
         data = [json.loads(line) for line in f if line.strip()]
 
     results = []
-    ifd_values = []
+    all_ifds = []
 
-    with torch.inference_mode():
+    for item in tqdm.tqdm(data, desc="Computing IFD", unit="sample"):
-        for item in tqdm.tqdm(data, desc="Computing IFD", unit="sample"):
+        if data_format == "messages":
-            instruction = item[instr_key]
+            scores = score_messages(
-            response = item[resp_key]
+                model, tokenizer, item.get("messages", []), device, max_len
-            scores = compute_ifd(
-                model,
-                tokenizer,
-                instruction,
-                response,
-                device,
-                max_len,
-                use_chat_template=use_chat_template,
             )
-            ifd_values.append(scores["ifd"])
+            if scores is None:
+                results.append({**item, "ifd": None, "ifd_turns": []})
+            else:
+                all_ifds.append(scores["ifd"])
+                results.append({**item, **scores})
+        else:
+            scores = score_plain(
+                model, tokenizer, item[instr_key], item[resp_key], device, max_len
+            )
+            all_ifds.append(scores["ifd"])
             results.append({**item, "ifd": scores["ifd"], "ifd_detail": scores})
 
     with open(output_file, "w", encoding="utf-8") as f:
         for item in results:
             f.write(json.dumps(item, ensure_ascii=False) + "\n")
 
-    valid_ifd = [v for v in ifd_values if v is not None]
+    valid_ifd = [v for v in all_ifds if v is not None]
     if valid_ifd:
-        import statistics
-
         print(f"\n{'=' * 50}")
         print(f"  Samples:    {len(data)}")
         print(f"  Valid IFD:  {len(valid_ifd)}")
@@ -266,29 +158,19 @@ def main():
     parser.add_argument("--param_path", type=str, required=True, help="Model directory")
     parser.add_argument("--input", type=str, required=True, help="Input JSONL file")
     parser.add_argument("--output", type=str, required=True, help="Output JSONL file")
+    parser.add_argument("--max_len", type=int, default=2048, help="Max token length")
     parser.add_argument(
-        "--instr_key",
+        "--format",
         type=str,
-        default="instruction",
+        default="plain",
-        help="Key for instruction field",
+        choices=["plain", "messages"],
+        help="Input format: 'plain' for instr_key+resp_key, 'messages' for messages array",
     )
     parser.add_argument(
-        "--resp_key",
+        "--instr_key", type=str, default="instruction", help="Key for instruction field"
-        type=str,
-        default="response",
-        help="Key for response field",
     )
     parser.add_argument(
-        "--max_len",
+        "--resp_key", type=str, default="response", help="Key for response field"
-        type=int,
-        default=2048,
-        help="Max token length (instruction truncated to fit)",
-    )
-    parser.add_argument(
-        "--no_chat_template",
-        action="store_true",
-        default=False,
-        help="Disable chat template, use raw text concatenation",
     )
     args = parser.parse_args()
 
@@ -299,7 +181,7 @@ def main():
         args.instr_key,
         args.resp_key,
         args.max_len,
-        use_chat_template=not args.no_chat_template,
+        data_format=args.format,
     )