assets/docs/preprocessing.md

@@ -268,7 +268,7 @@ When `sources` is set, `sections` is ignored.
 | `storage_format` | str | `"bin"` | `"bin"` (mmap) or `"h5"` |
 | `max_tokens_per_shard` | int | `100000000` | Flush threshold in cumulative tokens |
 | `dtype` | dict[str, str] | `{}` | Per-key tensor dtype override (e.g. `{"loss_mask": "bool"}`) |
-| `position_ids_mode` | str | `"doc_reset"` | How to compute position_ids: `"none"`, `"doc_reset"`, `"continuous"` |
+| `position_ids_mode` | str | `"none"` | How to compute position_ids: `"none"`, `"doc_reset"`, `"continuous"` |
 
 ---
 

astrai/config/preprocess_config.py

@@ -96,7 +96,7 @@ class OutputConfig(BaseConfig):
     storage_format: str = "bin"
     max_tokens_per_shard: int = 100_000_000
     dtype: Dict[str, str] = field(default_factory=dict)
-    position_ids_mode: str = "doc_reset"
+    position_ids_mode: str = "none"
 
 
 @dataclass

astrai/preprocessing/pipeline.py

@@ -144,7 +144,7 @@ class Pipeline:
         for key in list(bucket.keys()):
             if key in result:
                 continue
-            bucket[key].append([0] * len(ids))
+            bucket[key].append([1] * len(ids))
 
     def _iter_items(self):
         for path in self.paths:
@@ -160,12 +160,6 @@ class Pipeline:
             idx = shard_idx[domain]
 
             pp = self.config.preprocessing
-            original_sequences = keys.get("sequence", [])
-            mode = self.config.output.position_ids_mode
-
-            if mode == "doc_reset" and original_sequences:
-                keys["position_ids"] = [list(range(len(s))) for s in original_sequences]
-
             keys = self._packer.apply(dict(keys), pp.max_packed_len, pp.truncation_mode)
 
             tensors: Dict[str, List[torch.Tensor]] = {}
@@ -177,10 +171,9 @@ class Pipeline:
                     torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
                 ]
 
-            if mode == "continuous" and original_sequences:
+            pos_ids = self._position_id.generate(keys.get("sequence", []))
-                pos_ids = self._position_id.generate(keys.get("sequence", []))
+            if pos_ids:
-                if pos_ids:
+                tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
-                    tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
 
             self._writer.save(self.output_dir, domain, idx, tensors)
             shard_idx[domain] = idx + 1

scripts/eval/evaluate_ifd.py

@@ -1,14 +1,31 @@
 """IFD (Instruction Following Difficulty) data quality scoring.
 
-IFD = conditional_NLL / unconditional_NLL
+Computes IFD scores for instruction-response pairs to guide data selection.
+IFD = conditional_NLL / unconditional_NLL, where:
 
-- Messages format: plain text concatenation (no chat template)
+- conditional_NLL: average CE loss on response tokens given instruction context
-- Plain format: raw instr_key + resp_key fields
+- 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
 """
 
 import argparse
 import json
-import statistics
 
 import torch
 import torch.nn.functional as F
@@ -18,223 +35,217 @@ from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 
 
-def _pack_bins(pairs, max_len):
+def compute_ifd(
-    """BFD bin packing: pack (c+r) into bins of max total length."""
+    model,
-    indexed = sorted(enumerate(pairs), key=lambda x: -(len(x[1][0]) + len(x[1][1])))
+    tokenizer,
-    bins = []  # each bin: list of (orig_idx, ctx_ids, resp_ids)
+    instruction: str,
-    lengths = []
+    response: str,
-    for orig_idx, (c, r) in indexed:
+    device: str,
-        size = len(c) + len(r)
+    max_len: int = 2048,
-        best_bin = -1
+    use_chat_template: bool = False,
-        for bi, rem in enumerate(lengths):
+) -> dict:
-            if rem >= size:
+    if use_chat_template:
-                if best_bin < 0 or rem < lengths[best_bin]:
+        return _compute_ifd_with_template(
-                    best_bin = bi
+            model, tokenizer, instruction, response, device, max_len
-        if best_bin >= 0:
+        )
-            bins[best_bin].append((orig_idx, c, r))
+    return _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len)
-            lengths[best_bin] -= size
+
+
+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 {
+            "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:
-            bins.append([(orig_idx, c, r)])
+            instr_ids = instr_ids[overflow:]
-            lengths.append(max_len - size)
-    return bins
 
+    if not instr_ids:
+        return {
+            "L_cond": None,
+            "L_uncond": None,
+            "ifd": None,
+            "error": "response too long for context",
+        }
 
-@torch.inference_mode()
+    instr_len = len(instr_ids)
-def _score_batch(pairs, model, device, max_len=2048):
+    resp_len = len(resp_ids)
-    """BFD-packed IFD: pack items into bins, one forward pass per bin."""
-    if not pairs:
-        return []
-    bins = _pack_bins(pairs, max_len)
 
-    result = [None] * len(pairs)
+    qa_ids = instr_ids + resp_ids
 
-    for bin_items in bins:
+    with torch.inference_mode():
-        seq_ids = []
+        logits_qa = model(torch.tensor([qa_ids], device=device, dtype=torch.long))[
-        global_pos = []  # doc-reset position IDs for RoPE
+            "logits"
-        doc_ids = []  # document index for attention mask
+        ][0]
-        doc_offsets = []
+        logits_resp = model(torch.tensor([resp_ids], device=device, dtype=torch.long))[
-
-        for di, (orig_idx, c, r) in enumerate(bin_items):
-            ctx_len = len(c)
-            start = len(seq_ids)
-            item_len = len(c) + len(r)
-            seq_ids.extend(c)
-            seq_ids.extend(r)
-            end = len(seq_ids)
-            global_pos.extend(range(item_len))
-            doc_ids.extend([di] * item_len)
-            doc_offsets.append((start, end, orig_idx, ctx_len))
-
-        full_ids = torch.tensor([seq_ids], device=device, dtype=torch.long)
-        pos_ids = torch.tensor([global_pos], device=device, dtype=torch.long)
-        T = len(seq_ids)
-        causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=device))
-        doc_t = torch.tensor([doc_ids], device=device)
-        doc_mask = doc_t.unsqueeze(-1) == doc_t.unsqueeze(-2)
-        attn_mask = (causal & doc_mask[0]).unsqueeze(0).unsqueeze(0)
-        logits_full = model(full_ids, position_ids=pos_ids, input_mask=attn_mask)[
             "logits"
         ][0]
 
-        for start, end, orig_idx, ctx_len in doc_offsets:
+    resp_logits = logits_qa[instr_len - 1 : -1]
-            rl = end - start - ctx_len
+    resp_targets = logits_resp.new_tensor(resp_ids, dtype=torch.long)
-            if rl < 2:
+    L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()
-                continue
-            resp_start = start + ctx_len - 1
-            resp_logits = logits_full[resp_start : end - 1]
-            resp_targets = torch.tensor(
-                seq_ids[start + ctx_len : end], device=device, dtype=torch.long
-            )
-            L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()
-            result[orig_idx] = (L_cond, rl)
 
-    # unconditional pass: batch all responses separately (sorted by length)
+    unp_logits = logits_resp[:-1]
-    resp_seqs = [
+    unp_targets = logits_resp.new_tensor(resp_ids[1:], dtype=torch.long)
-        (i, result[i][1], pairs[i][1])
+    L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
-        for i in range(len(pairs))
-        if result[i] is not None
-    ]
-    if resp_seqs:
-        resp_seqs.sort(key=lambda x: -x[1])
-        r_batch = torch.zeros(
-            len(resp_seqs),
-            max(len(r) for _, _, r in resp_seqs),
-            dtype=torch.long,
-            device=device,
-        )
-        for ri, (_, rl, r_ids) in enumerate(resp_seqs):
-            r_batch[ri, :rl] = torch.tensor(r_ids, dtype=torch.long)
-        logits_resp = model(r_batch)["logits"]
 
-        for ri, (orig_idx, rl, _) in enumerate(resp_seqs):
+    ifd = L_cond / L_uncond if L_uncond > 0 else None
-            L_cond = result[orig_idx][0]
-            unp_logits = logits_resp[ri, : rl - 1]
-            unp_targets = r_batch[ri, 1:rl]
-            L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
-            ifd = L_cond / L_uncond if L_uncond > 0 else None
-            result[orig_idx] = {
-                "L_cond": round(L_cond, 6),
-                "L_uncond": round(L_uncond, 6),
-                "ifd": round(ifd, 6) if ifd is not None else None,
-                "resp_len": rl,
-            }
 
-    return result
-
-
-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_batch([(ctx_ids, resp_ids)], model, device, max_len)[0]
-
-
-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((ctx_ids, resp_ids))
-    if not turns:
-        return None
-    raw_scores = _score_batch(turns, model, device, max_len)
-    valid = [s for s in raw_scores if s is not None and s["ifd"] is not None]
-    if not valid:
-        return {"ifd": None, "ifd_turns": raw_scores}
-    avg = sum(s["ifd"] for s in valid) / len(valid)
     return {
-        "ifd": avg,
+        "L_cond": round(L_cond, 6),
-        "ifd_detail": valid[0] if len(valid) == 1 else None,
+        "L_uncond": round(L_uncond, 6),
-        "ifd_turns": raw_scores,
+        "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 process_file(
-    param_path,
+    param_path: str,
-    input_file,
+    input_file: str,
-    output_file,
+    output_file: str,
-    instr_key,
+    instr_key: str,
-    resp_key,
+    resp_key: str,
-    max_len=2048,
+    max_len: int = 2048,
-    data_format="plain",
+    use_chat_template: bool = False,
-    batch_size=1,
-    device=None,
 ):
-    if device is None:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
-    dtype = torch.bfloat16 if "cuda" in device else torch.float32
 
     model = AutoModel.from_pretrained(param_path)
     tokenizer = AutoTokenizer.from_pretrained(param_path)
     model.to(device=device, dtype=dtype)
     model.eval()
 
-    with open(input_file, encoding="utf-8") as f:
+    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:
         data = [json.loads(line) for line in f if line.strip()]
 
     results = []
-    all_ifds = []
+    ifd_values = []
-    buffer = []
 
-    for item in tqdm.tqdm(data, desc="Computing IFD", unit="sample"):
+    with torch.inference_mode():
-        if data_format == "messages":
+        for item in tqdm.tqdm(data, desc="Computing IFD", unit="sample"):
-            turns = []
+            instruction = item[instr_key]
-            for i, msg in enumerate(item.get("messages", [])):
+            response = item[resp_key]
-                if msg.get("role") != "assistant":
+            scores = compute_ifd(
-                    continue
+                model,
-                ctx_text = "\n\n".join(m["content"] for m in item["messages"][:i])
+                tokenizer,
-                ctx_ids = tokenizer.encode(ctx_text)
+                instruction,
-                resp_ids = tokenizer.encode(msg["content"], add_special_tokens=False)
+                response,
-                ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
+                device,
-                if ctx_ids and resp_ids:
+                max_len,
-                    turns.append((ctx_ids, resp_ids))
+                use_chat_template=use_chat_template,
-            if not turns:
+            )
-                results.append({**item, "ifd": None, "ifd_turns": []})
+            ifd_values.append(scores["ifd"])
-                continue
+            results.append({**item, "ifd": scores["ifd"], "ifd_detail": scores})
-            buffer.append((item, turns, "messages"))
-        else:
-            ctx_ids = tokenizer.encode(item[instr_key], add_special_tokens=False)
-            resp_ids = tokenizer.encode(item[resp_key], add_special_tokens=False)
-            ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
-            if not ctx_ids or not resp_ids:
-                results.append({**item, "ifd": None, "ifd_detail": {"error": "empty"}})
-                continue
-            buffer.append((item, [(ctx_ids, resp_ids)], "plain"))
-
-        if len(buffer) >= batch_size:
-            _flush_buffer(buffer, results, all_ifds, model, device, max_len)
-
-    if buffer:
-        _flush_buffer(buffer, results, all_ifds, model, device, max_len)
 
     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 all_ifds if v is not None]
+    valid_ifd = [v for v in ifd_values 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)}")
@@ -248,41 +259,6 @@ def process_file(
     print(f"Results saved to {output_file}")
 
 
-def _flush_buffer(buffer, results, all_ifds, model, device, max_len=2048):
-    all_pairs = []
-    indices = []
-    for item, turns, fmt in buffer:
-        start = len(all_pairs)
-        all_pairs.extend(turns)
-        indices.append((item, turns, fmt, start, len(all_pairs)))
-
-    raw = _score_batch(all_pairs, model, device, max_len)
-
-    for item, turns, fmt, start, end in indices:
-        turn_scores = raw[start:end]
-        if fmt == "messages":
-            valid = [s for s in turn_scores if s is not None and s["ifd"] is not None]
-            if not valid:
-                results.append({**item, "ifd": None, "ifd_turns": turn_scores})
-            else:
-                avg = sum(s["ifd"] for s in valid) / len(valid)
-                all_ifds.append(avg)
-                results.append(
-                    {
-                        **item,
-                        "ifd": avg,
-                        "ifd_detail": valid[0] if len(valid) == 1 else None,
-                        "ifd_turns": turn_scores,
-                    }
-                )
-        else:
-            score = turn_scores[0]
-            all_ifds.append(score["ifd"])
-            results.append({**item, "ifd": score["ifd"], "ifd_detail": score})
-
-    buffer.clear()
-
-
 def main():
     parser = argparse.ArgumentParser(
         description="Compute IFD scores for instruction-response data"
@@ -290,24 +266,30 @@ 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(
-        "--format",
+        "--instr_key",
         type=str,
-        default="plain",
+        default="instruction",
-        choices=["plain", "messages"],
+        help="Key for instruction field",
-        help="Input format",
     )
     parser.add_argument(
-        "--instr_key", type=str, default="instruction", help="Key for instruction field"
+        "--resp_key",
+        type=str,
+        default="response",
+        help="Key for response field",
     )
     parser.add_argument(
-        "--resp_key", type=str, default="response", help="Key for response field"
+        "--max_len",
+        type=int,
+        default=2048,
+        help="Max token length (instruction truncated to fit)",
     )
     parser.add_argument(
-        "--batch_size", type=int, default=8, help="Batch size for model forward passes"
+        "--no_chat_template",
+        action="store_true",
+        default=False,
+        help="Disable chat template, use raw text concatenation",
     )
-    parser.add_argument("--device", type=str, default=None, help="Device (e.g. cuda:0)")
     args = parser.parse_args()
 
     process_file(
@@ -317,9 +299,7 @@ def main():
         args.instr_key,
         args.resp_key,
         args.max_len,
-        data_format=args.format,
+        use_chat_template=not args.no_chat_template,
-        batch_size=args.batch_size,
-        device=args.device,
     )