diff --git a/scripts/eval/evaluate_ifd.py b/scripts/eval/evaluate_ifd.py index 5bb5843..bb84c7d 100644 --- a/scripts/eval/evaluate_ifd.py +++ b/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, where: +IFD = conditional_NLL / unconditional_NLL -- conditional_NLL: average CE loss on response tokens given instruction context -- unconditional_NLL: average CE loss on response tokens alone - -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 +- Messages format: plain text concatenation (no chat template) +- Plain format: raw instr_key + resp_key fields """ 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( - model, - 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] - +def _score(context_ids, resp_ids, model, device): + """Core IFD computation: context → L_cond, response alone → L_uncond.""" if not resp_ids: - return { - "L_cond": None, - "L_uncond": None, - "ifd": None, - "error": "empty response", - } - - qa_len = len(instr_ids) + len(resp_ids) - if qa_len > max_len: - 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) + return None + full_ids = context_ids + resp_ids + inp_full = torch.tensor([full_ids], device=device, dtype=torch.long) + inp_resp = torch.tensor([resp_ids], device=device, dtype=torch.long) + logits_full = model(inp_full)["logits"][0] + logits_resp = model(inp_resp)["logits"][0] + ctx_len = len(context_ids) + resp_logits = logits_full[ctx_len - 1 : -1] + resp_targets = torch.tensor(resp_ids, device=device, 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, - input_file: str, - output_file: str, - instr_key: str, - resp_key: str, - max_len: int = 2048, - use_chat_template: bool = False, + param_path, + input_file, + output_file, + instr_key, + resp_key, + max_len=2048, + 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: - 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: + with open(input_file, 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"): - instruction = item[instr_key] - response = item[resp_key] - scores = compute_ifd( - model, - tokenizer, - instruction, - response, - device, - max_len, - use_chat_template=use_chat_template, + for item in tqdm.tqdm(data, desc="Computing IFD", unit="sample"): + if data_format == "messages": + scores = score_messages( + model, tokenizer, item.get("messages", []), device, max_len ) - 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", - help="Key for instruction field", + default="plain", + choices=["plain", "messages"], + help="Input format: 'plain' for instr_key+resp_key, 'messages' for messages array", ) parser.add_argument( - "--resp_key", - type=str, - default="response", - help="Key for response field", + "--instr_key", type=str, default="instruction", help="Key for instruction field" ) parser.add_argument( - "--max_len", - 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", + "--resp_key", type=str, default="response", help="Key for response field" ) 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, )