assets/docs/architecture.md

@@ -363,16 +363,6 @@ classDiagram
         class TextMaskBuilder {
             +build(item, config, tokenizer) Optional[dict]
         }
-
-        class Pipeline {
-            +PipelineConfig config
-            +List[str] paths
-            +str output_dir
-            +str tokenizer_path
-            +BaseMaskBuilder mask_builder
-            +transform(item) Optional[dict]
-            +run()
-        }
     }
 
     namespace tokenize {
@@ -1102,8 +1092,6 @@ classDiagram
     KvcacheView o-- Storage
     SamplingPipeline o-- BaseSamplingStrategy
     BaseDataset o-- Store
-    Pipeline o-- PipelineConfig
-    Pipeline o-- BaseMaskBuilder
 
     %% --- Dependency (uses temporarily) ---
     TrainConfig ..> BaseStrategy : selects

assets/docs/preprocessing.md

@@ -186,8 +186,6 @@ Pure tokenization. No `loss_mask` is produced. Used for pretraining.
 
 ## Output Layout
 
-### Single-Shard (`bin`)
-
 ```
 output_dir/
   __default__/              # when domain_key is null
@@ -200,59 +198,6 @@ output_dir/
     loss_mask.bin
 ```
 
-### Multi-Shard (`bin`)
-
-When `max_tokens_per_shard` is exceeded, bin output is split into numbered shard subdirectories:
-
-```
-output_dir/
-  __default__/
-    shard_0000/
-      meta.json
-      sequence.bin
-      loss_mask.bin
-    shard_0001/
-      meta.json
-      sequence.bin
-      loss_mask.bin
-```
-
-`MmapStore` automatically discovers and merges all shards under the domain directory.
-
-### H5 Output
-
-HDF5 files are always named with a shard index, avoiding overwrite regardless of `max_tokens_per_shard`:
-
-```
-output_dir/
-  __default__/
-    data_0000.h5            # each H5 contains key→dataset groups
-    data_0001.h5
-  wiki/
-    data_0000.h5
-```
-
-## Python API Usage
-
-```python
-from astrai.preprocessing.pipeline import Pipeline
-from astrai.config.preprocess_config import PipelineConfig
-
-config = PipelineConfig.from_json("sft_pipeline.json")
-Pipeline(
-    config,
-    ["data_part1.jsonl", "data_part2.jsonl"],
-    output_dir="output/",
-    tokenizer_path="params"
-).run()
-```
-
-Or from the CLI:
-
-```bash
-python scripts/tools/preprocess.py data/*.jsonl -o output/ -c sft.json
-```
-
 ## Extension
 
 Register a custom builder for new formats:

astrai/config/preprocess_config.py

@@ -1,14 +1,20 @@
 """Pipeline configuration for JSONL preprocessing."""
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional
+from typing import Dict, Optional
 
 from astrai.config.base import BaseConfig
 
 
 @dataclass
 class InputConfig(BaseConfig):
-    sections: Optional[List[Dict]] = None
+    type: str = "chat"
+    messages_key: str = "messages"
+    prompt_key: str = "prompt"
+    response_key: str = "response"
+    text_key: str = "text"
 
 
 @dataclass
@@ -16,6 +22,7 @@ class ProcessingConfig(BaseConfig):
     max_seq_len: int = 2048
     min_chars: int = 50
     max_chars: int = 2_000_000
+    deduplicate: bool = True
     max_items: Optional[int] = None
 
 
@@ -24,7 +31,6 @@ class OutputConfig(BaseConfig):
     domain_key: Optional[str] = None
     storage_format: str = "bin"
     max_tokens_per_shard: int = 100_000_000
-    dtype: Dict[str, str] = field(default_factory=dict)
 
 
 @dataclass

astrai/dataset/storage.py

@@ -117,11 +117,7 @@ def detect_format(load_path: str) -> str:
     if h5_files:
         return "h5"
     bin_files = list(root.rglob("*.bin"))
-    if bin_files:
+    if bin_files and (root / "meta.json").exists():
-        has_meta = (root / "meta.json").exists() or len(
-            list(root.rglob("meta.json"))
-        ) > 0
-        if has_meta:
         return "bin"
     raise FileNotFoundError(f"No supported data files found at {load_path}")
 
@@ -248,17 +244,7 @@ class MmapStore(Store):
 
     def load(self, path: str):
         self._mmap_refs = []
-        root = Path(path)
+        raw = load_bin(path)
-        all_raw: Dict[str, List[Tensor]] = {}
+        self._normalize(raw)
-        meta_paths = list(root.rglob("meta.json"))
-        for meta_path in meta_paths:
-            raw = load_bin(str(meta_path.parent))
-            for key, tensors in raw.items():
-                if key not in all_raw:
-                    all_raw[key] = []
-                all_raw[key].extend(tensors)
-        if not meta_paths:
-            raise FileNotFoundError(f"No meta.json found under {path}")
-        self._normalize(all_raw)
         for tensors in self._data.values():
             self._mmap_refs.extend(tensors)

astrai/inference/api/anthropic.py

@@ -1,6 +1,5 @@
 """Anthropic message completion response builder."""
 
-import time
 import uuid
 from typing import Any, Dict, List, Tuple, Union
 
@@ -40,7 +39,7 @@ class AnthropicResponseBuilder(ResponseBuilder):
         prompt = engine.tokenizer.apply_chat_template(messages, tokenize=False)
         ctx = GenContext(
             resp_id=f"msg_{uuid.uuid4().hex[:24]}",
-            created=int(time.time()),
+            created=0,
             model=request.model,
             prompt_tokens=0,
         )

astrai/inference/api/openai.py

@@ -1,7 +1,5 @@
 """OpenAI chat completion response builder."""
 
-import logging
-import time
 import uuid
 from typing import Any, Dict, List, Tuple
 
@@ -15,16 +13,6 @@ from astrai.inference.api.protocol import (
 )
 from astrai.inference.engine import InferenceEngine
 
-logger = logging.getLogger(__name__)
-
-_UNSUPPORTED_PARAMS = (
-    "n",
-    "presence_penalty",
-    "frequency_penalty",
-    "logit_bias",
-    "user",
-)
-
 
 class OpenAIResponseBuilder(ResponseBuilder):
     def prepare(
@@ -36,26 +24,9 @@ class OpenAIResponseBuilder(ResponseBuilder):
         self._resp_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
         self._model = request.model
 
-        for param in _UNSUPPORTED_PARAMS:
-            value = getattr(request, param, None)
-            fields = getattr(type(request), "model_fields", {})
-            default = fields[param].default if param in fields else None
-            if value is not None and value != default:
-                logger.warning(
-                    "ChatCompletionRequest param '%s'=%r is not supported and will be ignored",
-                    param,
-                    value,
-                )
-            if value is not None and value != default:
-                logger.warning(
-                    "ChatCompletionRequest param '%s'=%r is not supported and will be ignored",
-                    param,
-                    value,
-                )
-
         ctx = GenContext(
             resp_id=self._resp_id,
-            created=int(time.time()),
+            created=0,
             model=self._model,
             prompt_tokens=0,
         )

astrai/inference/api/protocol.py

@@ -64,7 +64,7 @@ class StopChecker:
 class ResponseBuilder(ABC):
     """Interface for protocol-specific response formatting.
 
-    A new protocol requires one concrete builder implementing 5 methods.
+    A new protocol requires one concrete builder implementing 6 methods.
     """
 
     @abstractmethod

astrai/preprocessing/__init__.py

@@ -1,14 +1,19 @@
 from astrai.preprocessing.builder import (
     BaseMaskBuilder,
+    ChatMaskBuilder,
+    InstructionMaskBuilder,
     MaskBuilderFactory,
-    SectionedMaskBuilder,
+    TextMaskBuilder,
 )
-from astrai.preprocessing.pipeline import Pipeline, filter_by_length
+from astrai.preprocessing.pipeline import Pipeline, dedup_signature, filter_by_length
 
 __all__ = [
     "BaseMaskBuilder",
+    "ChatMaskBuilder",
+    "InstructionMaskBuilder",
     "MaskBuilderFactory",
-    "SectionedMaskBuilder",
+    "TextMaskBuilder",
     "Pipeline",
+    "dedup_signature",
     "filter_by_length",
 ]

astrai/preprocessing/builder.py

@@ -1,11 +1,13 @@
 """Mask building strategies for preprocessing pipeline.
 
-The single :class:`SectionedMaskBuilder` handles all input formats
+Each builder knows how to tokenize one input format and construct
-via declarative ``input.sections`` config.
+the loss_mask according to declarative mask rules from the config.
 """
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import List, Optional
 
 from astrai.factory import BaseFactory
 
@@ -38,122 +40,122 @@ def _extract_domain(item: dict, domain_key: Optional[str]) -> str:
     return val if isinstance(val, str) else "__default__"
 
 
-def _resolve_action(action: str, role: str, config) -> str:
+@MaskBuilderFactory.register("chat")
-    """Resolve action to "train" or "mask".
+class ChatMaskBuilder(BaseMaskBuilder):
+    """Mask by role via message-level tokenisation with role-span tracking.
 
-    - ``"train"`` / ``"mask"`` → literal
+    For each message, renders the chat template for that single message,
-    - ``"$role"`` → look up ``role`` in ``config.mask``, fall back to ``config.mask_default``
+    encodes individually, and records its token span + role action.
-    """
+    The concatenated sequence receives a loss_mask built from span rules.
-    if action == "$role":
-        return config.mask.get(role, config.mask_default)
-    return action
-
-
-@MaskBuilderFactory.register("sectioned")
-class SectionedMaskBuilder(BaseMaskBuilder):
-    """Config-driven builder: iterates over ``input.sections`` in order.
-
-    Each section specifies a JSONL field + mask action.
-
-    Section spec::
-
-        {
-            "field":    "messages",   # JSONL key
-            "action":   "$role",      # "train" | "mask" | "$role"
-            "template": true,         # apply chat_template per message (optional)
-            "add_special_tokens": false  # override encode flag (optional)
-        }
-
-    Example configs::
-
-        # Chat
-        {"input": {"sections": [
-            {"field": "messages", "action": "$role", "template": true}
-        ]}}
-
-        # Instruction
-        {"input": {"sections": [
-            {"field": "prompt",   "action": "mask", "add_special_tokens": true},
-            {"field": "response", "action": "train"}
-        ]}}
-
-        # Text
-        {"input": {"sections": [
-            {"field": "text", "action": "train"}
-        ]}}
     """
 
     def build(self, item: dict, config, tokenizer) -> Optional[dict]:
-        sections = config.input.sections
+        messages = item.get(config.input.messages_key)
-        if not sections:
+        if not isinstance(messages, list) or not messages:
             return None
 
-        all_ids: list[int] = []
+        all_ids: List[int] = []
-        loss_mask: list[int] = []
+        spans: List[tuple] = []
 
-        has_template = any(s.get("template") for s in sections)
+        if tokenizer.bos_token_id is not None:
-        is_text_config = not has_template and all(
-            s["action"] == "train" for s in sections
-        )
-
-        if has_template and tokenizer.bos_token_id is not None:
             all_ids.append(tokenizer.bos_token_id)
-            loss_mask.append(0)
 
-        first_section = True
-        for sec in sections:
-            field = sec["field"]
-            action = sec["action"]
-            use_template = sec.get("template", False)
-            add_special = sec.get(
-                "add_special_tokens", not use_template and first_section
-            )
-
-            if use_template:
-                messages = item.get(field)
-                if not isinstance(messages, list) or not messages:
-                    continue
         for msg in messages:
             role = msg.get("role", "")
-                    act = _resolve_action(action, role, config)
+            action = config.mask.get(role, config.mask_default)
+
             rendered = tokenizer.apply_chat_template(
                 [msg], tokenize=False, add_generation_prompt=False
             )
             ids = tokenizer.encode(rendered, add_special_tokens=False)
-                    all_ids.extend(ids)
-                    val = 1 if act == "train" else 0
-                    loss_mask.extend([val] * len(ids))
-            else:
-                text = str(item.get(field, ""))
-                if not text.strip():
-                    continue
-                if is_text_config:
-                    pp = config.preprocessing
-                    if pp.min_chars > 0 and len(text) < pp.min_chars:
-                        continue
-                    if len(text) > pp.max_chars:
-                        continue
-                ids = tokenizer.encode(text, add_special_tokens=add_special)
-                all_ids.extend(ids)
-                val = 1 if action == "train" else 0
-                loss_mask.extend([val] * len(ids))
 
-            first_section = False
+            start = len(all_ids)
+            all_ids.extend(ids)
+            spans.append((start, len(all_ids), action))
+
+        if len(all_ids) <= 1:
+            return None
 
         max_len = config.preprocessing.max_seq_len
         all_ids = all_ids[:max_len]
-        loss_mask = loss_mask[: len(all_ids)]
 
-        if not all_ids:
+        loss_mask = [0] * len(all_ids)
-            return None
+        for start, end, action in spans:
+            if start >= len(all_ids):
+                break
+            e = min(end, len(all_ids))
+            if action == "train":
+                loss_mask[start:e] = [1] * (e - start)
 
-        if has_template and len(all_ids) <= 1:
+        return {
-            return None
+            "ids": all_ids,
-
+            "loss_mask": loss_mask,
-        result: dict = {
+            "domain": _extract_domain(item, config.output.domain_key),
-            "sequence": all_ids,
+        }
+
+
+@MaskBuilderFactory.register("instruction")
+class InstructionMaskBuilder(BaseMaskBuilder):
+    """Mask by prompt / response field boundary.
+
+    Encodes prompt and response independently, then fills mask
+    according to ``prompt`` / ``response`` entries in the mask config.
+    """
+
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        prompt = str(item.get(config.input.prompt_key, ""))
+        response = str(item.get(config.input.response_key, ""))
+
+        if not prompt.strip() and not response.strip():
+            return None
+
+        prompt_ids = tokenizer.encode(prompt, add_special_tokens=True)
+        response_ids = tokenizer.encode(response, add_special_tokens=False)
+
+        max_len = config.preprocessing.max_seq_len
+        full_ids = (prompt_ids + response_ids)[:max_len]
+
+        prompt_action = config.mask.get("prompt", config.mask_default)
+        response_action = config.mask.get("response", config.mask_default)
+
+        p_len = min(len(prompt_ids), len(full_ids))
+        r_len = len(full_ids) - p_len
+
+        loss_mask = []
+        if prompt_action == "train":
+            loss_mask += [1] * p_len
+        else:
+            loss_mask += [0] * p_len
+
+        if response_action == "train":
+            loss_mask += [1] * r_len
+        else:
+            loss_mask += [0] * r_len
+
+        return {
+            "ids": full_ids,
+            "loss_mask": loss_mask,
+            "domain": _extract_domain(item, config.output.domain_key),
+        }
+
+
+@MaskBuilderFactory.register("text")
+class TextMaskBuilder(BaseMaskBuilder):
+    """Plain tokenisation — no mask, used for pre-training data."""
+
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        text = item.get(config.input.text_key, "")
+        if not isinstance(text, str) or not text.strip():
+            return None
+
+        pp = config.preprocessing
+        if not (pp.min_chars <= len(text) <= pp.max_chars):
+            return None
+
+        ids = tokenizer.encode(text, add_special_tokens=True)
+        ids = ids[: pp.max_seq_len]
+
+        return {
+            "ids": ids,
             "domain": _extract_domain(item, config.output.domain_key),
         }
-        if not all(m == 1 for m in loss_mask):
-            result["loss_mask"] = loss_mask
-        return result

astrai/preprocessing/pipeline.py

@@ -1,40 +1,35 @@
 """Config-driven JSONL preprocessing pipeline.
 
 Composes a :class:`BaseMaskBuilder` (selected by ``input.type``) with
-sharding and flush to ``.h5`` / ``.bin`` storage.
+deduplication, sharding, and flush to ``.h5`` / ``.bin`` storage.
 """
 
+from __future__ import annotations
+
+import hashlib
 import json
 import os
 from collections import defaultdict
-from itertools import chain
+from typing import List, Optional
-from typing import Optional
 
 import torch
 import tqdm
 
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.dataset.storage import save_bin, save_h5
-from astrai.preprocessing.builder import SectionedMaskBuilder
+from astrai.preprocessing.builder import MaskBuilderFactory
 from astrai.tokenize import AutoTokenizer
 
-_STR_TO_DTYPE: dict[str, torch.dtype] = {
-    "bool": torch.bool,
-    "uint8": torch.uint8,
-    "int8": torch.int8,
-    "int16": torch.int16,
-    "int32": torch.int32,
-    "int64": torch.int64,
-    "float16": torch.float16,
-    "float32": torch.float32,
-    "float64": torch.float64,
-}
-
 
 def filter_by_length(text: str, min_len: int = 50, max_len: int = 2_000_000) -> bool:
     return min_len <= len(text) <= max_len
 
 
+def dedup_signature(item: dict) -> str:
+    raw = json.dumps(item, sort_keys=True, ensure_ascii=False)
+    return hashlib.md5(raw[:200].encode()).hexdigest()
+
+
 class Pipeline:
     """Tokenization pipeline driven by a declarative :class:`PipelineConfig`.
 
@@ -47,7 +42,7 @@ class Pipeline:
     def __init__(
         self,
         config: PipelineConfig,
-        input_paths: list[str],
+        input_paths: List[str],
         output_dir: str,
         tokenizer_path: str,
     ):
@@ -57,13 +52,15 @@ class Pipeline:
         self.output_dir = output_dir
         self.tokenizer_path = tokenizer_path
 
-        self.mask_builder = SectionedMaskBuilder()
+        self.mask_builder = MaskBuilderFactory.create(config.input.type)
 
     def transform(self, item: dict) -> Optional[dict]:
         return self.mask_builder.build(item, self.config, self._tokenizer)
 
     def run(self):
         self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path)
+
+        seen: set = set()
         domains: dict = defaultdict(lambda: defaultdict(list))
         total_tokens = 0
         shard_idx: dict[str, int] = defaultdict(int)
@@ -77,23 +74,24 @@ class Pipeline:
             if pp.max_items and count >= pp.max_items:
                 break
 
+            if pp.deduplicate:
+                sig = dedup_signature(item)
+                if sig in seen:
+                    continue
+                seen.add(sig)
+
             result = self.transform(item)
             if result is None:
                 continue
 
-            ids = result.pop("sequence")
+            ids = result["ids"]
             if not ids:
                 continue
 
-            domain = result.pop("domain", "__default__")
+            domain = result.get("domain", "__default__")
-            result["sequence"] = ids
+            domains[domain]["sequence"].append(ids)
-
+            if "loss_mask" in result:
-            bucket = domains[domain]
+                domains[domain]["loss_mask"].append(result["loss_mask"])
-            for key in list(bucket.keys()):
-                if key not in result:
-                    bucket[key].append([1] * len(ids))
-            for key, val in result.items():
-                bucket[key].append(val)
 
             count += 1
             total_tokens += len(ids)
@@ -122,16 +120,11 @@ class Pipeline:
             idx = shard_idx[domain]
             tensors = {}
             for key, ids_list in keys.items():
-                dt = _STR_TO_DTYPE.get(
+                tensors[key] = [torch.tensor(sum(ids_list, []), dtype=torch.long)]
-                    self.config.output.dtype.get(key, "int32"), torch.int32
-                )
-                tensors[key] = [
-                    torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
-                ]
             chunk_dir = os.path.join(self.output_dir, domain)
             fmt = self.config.output.storage_format
             if fmt == "bin":
-                save_bin(os.path.join(chunk_dir, f"shard_{idx:04d}"), tensors)
+                save_bin(chunk_dir, tensors)
             else:
                 save_h5(chunk_dir, f"data_{idx:04d}", tensors)
             shard_idx[domain] = idx + 1

astrai/tokenize/chat_template.py

@@ -1,10 +1,13 @@
+from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 
 from jinja2 import Template
 
+# Message type for chat messages
 type MessageType = Dict[str, Any]
 
 
+@dataclass
 class ChatTemplate:
     """A chat template with Jinja2 rendering support.
 
@@ -12,24 +15,23 @@ class ChatTemplate:
         name: Unique identifier for the template.
         template_str: Jinja2 template string.
         description: Optional description.
-        default_variables: Optional dictionary of default variable values.
+        default_variables: Optional dictionary of default variable values
+            that will be passed to the template if not overridden during rendering.
         special_tokens: Optional dictionary mapping token names to their string values.
+            These tokens are automatically added to the template variables.
     """
 
-    def __init__(
+    name: str
-        self,
+    template_str: str
-        name: str = "",
+    description: str = ""
-        template_str: str = "",
+    default_variables: Dict[str, Any] = None
-        description: str = "",
+    special_tokens: Dict[str, str] = None
-        default_variables: Optional[Dict[str, Any]] = None,
+
-        special_tokens: Optional[Dict[str, str]] = None,
+    def __post_init__(self):
-    ):
+        if self.default_variables is None:
-        self.name = name
+            self.default_variables = {}
-        self.template_str = template_str
+        if self.special_tokens is None:
-        self.description = description
+            self.special_tokens = {}
-        self.default_variables = default_variables or {}
-        self.special_tokens = special_tokens or {}
-        self._compiled: Template = Template(template_str)
 
     @classmethod
     def from_string(
@@ -41,7 +43,7 @@ class ChatTemplate:
     ) -> "ChatTemplate":
         """Create a ChatTemplate instance directly from a template string."""
         return cls(
-            name="",
+            name="",  # empty name for ad‑hoc templates
             template_str=template_str,
             description=description,
             default_variables=default_variables,
@@ -71,4 +73,5 @@ class ChatTemplate:
         if system_prompt is not None:
             variables["system_prompt"] = system_prompt
 
-        return self._compiled.render(**variables)
+        jinja_template = Template(self.template_str)
+        return jinja_template.render(**variables)

tests/data/test_preprocess.py

@@ -12,22 +12,22 @@ from astrai.config.preprocess_config import (
     ProcessingConfig,
 )
 from astrai.preprocessing.builder import (
+    ChatMaskBuilder,
+    InstructionMaskBuilder,
     MaskBuilderFactory,
-    SectionedMaskBuilder,
+    TextMaskBuilder,
 )
-from astrai.preprocessing.pipeline import Pipeline, filter_by_length
+from astrai.preprocessing.pipeline import Pipeline, dedup_signature, filter_by_length
 from astrai.tokenize import AutoTokenizer
 
-_SPECIAL_TOKENS_CONFIG = {
+_SPECIAL_TOKENS = [
-    "bos_token": "<|begin_of_sentence|>",
+    "<unk>",
-    "eos_token": "<|end_of_sentence|>",
+    "<pad>",
-    "pad_token": "<|_pad_|>",
+    "<|begin_of_sentence|>",
-    "unk_token": "<|_unk_|>",
+    "<|end_of_sentence|>",
-    "im_start": "<|im_start|>",
+    "<|im_start|>",
-    "im_end": "<|im_end|>",
+    "<|im_end|>",
-}
+]
-
-_SPECIAL_TOKENS = list(_SPECIAL_TOKENS_CONFIG.values())
 
 _CHAT_TEMPLATE = (
     "{% for message in messages %}"
@@ -75,8 +75,8 @@ def _build_chat_tokenizer() -> AutoTokenizer:
     auto_tok._special_token_map = {
         "bos_token": "<|begin_of_sentence|>",
         "eos_token": "<|end_of_sentence|>",
-        "pad_token": "<|_pad_|>",
+        "pad_token": "<pad>",
-        "unk_token": "<|_unk_|>",
+        "unk_token": "<unk>",
     }
     auto_tok.set_chat_template(_CHAT_TEMPLATE)
     return auto_tok
@@ -96,19 +96,9 @@ def temp_dir():
     shutil.rmtree(d, ignore_errors=True)
 
 
-_CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}]
-
-_INSTRUCTION_SECTIONS = [
-    {"field": "prompt", "action": "mask", "add_special_tokens": True},
-    {"field": "response", "action": "train"},
-]
-
-_TEXT_SECTIONS = [{"field": "text", "action": "train"}]
-
-
 def make_chat_config():
     return PipelineConfig(
-        input=InputConfig(sections=_CHAT_SECTIONS),
+        input=InputConfig(type="chat", messages_key="messages"),
         mask={"system": "mask", "user": "mask", "assistant": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -117,7 +107,9 @@ def make_chat_config():
 
 def make_instruction_config():
     return PipelineConfig(
-        input=InputConfig(sections=_INSTRUCTION_SECTIONS),
+        input=InputConfig(
+            type="instruction", prompt_key="prompt", response_key="response"
+        ),
         mask={"prompt": "mask", "response": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -126,7 +118,7 @@ def make_instruction_config():
 
 def make_text_config():
     return PipelineConfig(
-        input=InputConfig(sections=_TEXT_SECTIONS),
+        input=InputConfig(type="text", text_key="text"),
         preprocessing=ProcessingConfig(
             max_seq_len=2048, min_chars=1, max_chars=2_000_000
         ),
@@ -137,59 +129,58 @@ class TestPipelineConfig:
     def test_default_values(self):
         config = PipelineConfig()
         assert config.version == 1
+        assert config.input.type == "chat"
         assert config.mask == {}
         assert config.mask_default == "mask"
         assert config.preprocessing.max_seq_len == 2048
         assert config.output.storage_format == "bin"
-        assert config.input.sections is None
 
     def test_from_dict_flat(self):
         data = {
             "version": 1,
-            "input": {
+            "input": {"type": "chat", "messages_key": "msgs"},
-                "sections": [{"field": "messages", "action": "$role", "template": True}]
-            },
             "mask": {"system": "mask", "assistant": "train"},
             "mask_default": "mask",
             "preprocessing": {"max_seq_len": 1024},
             "output": {"storage_format": "h5"},
         }
         config = PipelineConfig.from_dict(data)
-        assert config.input.sections == [
+        assert config.input.type == "chat"
-            {"field": "messages", "action": "$role", "template": True}
+        assert config.input.messages_key == "msgs"
-        ]
         assert config.mask == {"system": "mask", "assistant": "train"}
         assert config.preprocessing.max_seq_len == 1024
         assert config.output.storage_format == "h5"
 
     def test_to_dict_roundtrip(self):
         config = PipelineConfig(
-            input=InputConfig(sections=_INSTRUCTION_SECTIONS),
+            input=InputConfig(type="instruction", prompt_key="q", response_key="a"),
             mask={"prompt": "mask", "response": "train"},
             mask_default="mask",
         )
         d = config.to_dict()
         config2 = PipelineConfig.from_dict(d)
-        assert config2.input.sections == _INSTRUCTION_SECTIONS
+        assert config2.input.type == "instruction"
+        assert config2.input.prompt_key == "q"
         assert config2.mask == {"prompt": "mask", "response": "train"}
 
     def test_to_json_from_json(self, temp_dir):
         config = PipelineConfig(
-            input=InputConfig(sections=_TEXT_SECTIONS),
+            input=InputConfig(type="text", text_key="body"),
             mask={"text": "train"},
             mask_default="mask",
         )
         path = os.path.join(temp_dir, "config.json")
         config.to_json(path)
         loaded = PipelineConfig.from_json(path)
-        assert loaded.input.sections == _TEXT_SECTIONS
+        assert loaded.input.type == "text"
+        assert loaded.input.text_key == "body"
         assert loaded.mask == {"text": "train"}
 
 
 class TestChatMaskBuilder:
     def test_simple_chat_mask(self, chat_tokenizer):
         config = make_chat_config()
-        builder = SectionedMaskBuilder()
+        builder = ChatMaskBuilder()
         item = {
             "messages": [
                 {"role": "system", "content": "You are helpful."},
@@ -199,23 +190,23 @@ class TestChatMaskBuilder:
         }
         result = builder.build(item, config, chat_tokenizer)
         assert result is not None
-        assert "sequence" in result
+        assert "ids" in result
         assert "loss_mask" in result
-        assert len(result["sequence"]) == len(result["loss_mask"])
+        assert len(result["ids"]) == len(result["loss_mask"])
 
-        ids = chat_tokenizer.decode(result["sequence"], skip_special_tokens=False)
+        ids = chat_tokenizer.decode(result["ids"], skip_special_tokens=False)
 
         assert "system" in ids.lower() or "<|im_start|>system" in ids
         assert "assistant" in ids.lower() or "<|im_start|>assistant" in ids
 
-        total = len(result["sequence"])
+        total = len(result["ids"])
         trained = sum(result["loss_mask"])
         assert trained > 0, "At least assistant tokens should be trained"
         assert trained < total, "System and user tokens should be masked"
 
     def test_mask_only_assistant_trained(self, chat_tokenizer):
         config = make_chat_config()
-        builder = SectionedMaskBuilder()
+        builder = ChatMaskBuilder()
         item = {
             "messages": [
                 {"role": "user", "content": "What is 2+2?"},
@@ -224,7 +215,7 @@ class TestChatMaskBuilder:
         }
         result = builder.build(item, config, chat_tokenizer)
         mask = result["loss_mask"]
-        ids = result["sequence"]
+        ids = result["ids"]
 
         assert len(ids) == len(mask)
 
@@ -236,12 +227,12 @@ class TestChatMaskBuilder:
 
     def test_chat_all_masked(self, chat_tokenizer):
         config = PipelineConfig(
-            input=InputConfig(sections=_CHAT_SECTIONS),
+            input=InputConfig(type="chat", messages_key="messages"),
             mask={"system": "mask", "user": "mask", "assistant": "mask"},
             mask_default="mask",
             preprocessing=ProcessingConfig(max_seq_len=2048),
         )
-        builder = SectionedMaskBuilder()
+        builder = ChatMaskBuilder()
         item = {
             "messages": [
                 {"role": "system", "content": "You are helpful."},
@@ -253,12 +244,12 @@ class TestChatMaskBuilder:
 
     def test_chat_all_trained(self, chat_tokenizer):
         config = PipelineConfig(
-            input=InputConfig(sections=_CHAT_SECTIONS),
+            input=InputConfig(type="chat", messages_key="messages"),
             mask={},
             mask_default="train",
             preprocessing=ProcessingConfig(max_seq_len=2048),
         )
-        builder = SectionedMaskBuilder()
+        builder = ChatMaskBuilder()
         item = {
             "messages": [
                 {"role": "system", "content": "You are helpful."},
@@ -266,23 +257,23 @@ class TestChatMaskBuilder:
             ]
         }
         result = builder.build(item, config, chat_tokenizer)
-        assert sum(result["loss_mask"]) == len(result["sequence"]) - 1
+        assert sum(result["loss_mask"]) == len(result["ids"]) - 1
 
     def test_empty_messages_returns_none(self, chat_tokenizer):
         config = make_chat_config()
-        builder = SectionedMaskBuilder()
+        builder = ChatMaskBuilder()
         assert builder.build({"messages": []}, config, chat_tokenizer) is None
         assert builder.build({}, config, chat_tokenizer) is None
 
     def test_domain_extraction(self, chat_tokenizer):
         config = PipelineConfig(
-            input=InputConfig(sections=_CHAT_SECTIONS),
+            input=InputConfig(type="chat", messages_key="messages"),
             mask={"assistant": "train"},
             mask_default="mask",
             preprocessing=ProcessingConfig(max_seq_len=2048),
             output=OutputConfig(domain_key="source"),
         )
-        builder = SectionedMaskBuilder()
+        builder = ChatMaskBuilder()
         item = {
             "messages": [
                 {"role": "user", "content": "Hi"},
@@ -295,12 +286,12 @@ class TestChatMaskBuilder:
 
     def test_truncation_to_max_len(self, chat_tokenizer):
         config = PipelineConfig(
-            input=InputConfig(sections=_CHAT_SECTIONS),
+            input=InputConfig(type="chat", messages_key="messages"),
             mask={"assistant": "train"},
             mask_default="mask",
             preprocessing=ProcessingConfig(max_seq_len=10),
         )
-        builder = SectionedMaskBuilder()
+        builder = ChatMaskBuilder()
         item = {
             "messages": [
                 {
@@ -311,26 +302,26 @@ class TestChatMaskBuilder:
             ]
         }
         result = builder.build(item, config, chat_tokenizer)
-        assert len(result["sequence"]) <= 10
+        assert len(result["ids"]) <= 10
-        assert len(result["loss_mask"]) == len(result["sequence"])
+        assert len(result["loss_mask"]) == len(result["ids"])
 
 
 class TestInstructionMaskBuilder:
     def test_basic_instruction_mask(self, test_tokenizer):
         config = make_instruction_config()
-        builder = SectionedMaskBuilder()
+        builder = InstructionMaskBuilder()
         item = {"prompt": "Translate to French: Hello", "response": "Bonjour"}
         result = builder.build(item, config, test_tokenizer)
         assert result is not None
-        assert len(result["sequence"]) == len(result["loss_mask"])
+        assert len(result["ids"]) == len(result["loss_mask"])
 
     def test_prompt_masked_response_trained(self, test_tokenizer):
         config = make_instruction_config()
-        builder = SectionedMaskBuilder()
+        builder = InstructionMaskBuilder()
         item = {"prompt": "hello", "response": "world"}
         result = builder.build(item, config, test_tokenizer)
         mask = result["loss_mask"]
-        ids = result["sequence"]
+        ids = result["ids"]
 
         prompt_ids = test_tokenizer.encode("hello", add_special_tokens=True)
         response_ids = test_tokenizer.encode("world", add_special_tokens=False)
@@ -344,22 +335,17 @@ class TestInstructionMaskBuilder:
     def test_train_on_prompt(self, test_tokenizer):
         config = PipelineConfig(
             input=InputConfig(
-                sections=[
+                type="instruction", prompt_key="prompt", response_key="response"
-                    {
-                        "field": "prompt",
-                        "action": "train",
-                        "add_special_tokens": True,
-                    },
-                    {"field": "response", "action": "mask"},
-                ]
             ),
+            mask={"prompt": "train", "response": "mask"},
+            mask_default="mask",
             preprocessing=ProcessingConfig(max_seq_len=2048),
         )
-        builder = SectionedMaskBuilder()
+        builder = InstructionMaskBuilder()
         item = {"prompt": "hello", "response": "world"}
         result = builder.build(item, config, test_tokenizer)
         mask = result["loss_mask"]
-        ids = result["sequence"]
+        ids = result["ids"]
 
         prompt_ids = test_tokenizer.encode("hello", add_special_tokens=True)
         p_len = min(len(prompt_ids), len(ids))
@@ -369,37 +355,37 @@ class TestInstructionMaskBuilder:
 class TestTextMaskBuilder:
     def test_basic_text(self, test_tokenizer):
         config = make_text_config()
-        builder = SectionedMaskBuilder()
+        builder = TextMaskBuilder()
         item = {"text": "Hello world. This is a test document."}
         result = builder.build(item, config, test_tokenizer)
         assert result is not None
-        assert "sequence" in result
+        assert "ids" in result
-        assert len(result["sequence"]) > 0
+        assert len(result["ids"]) > 0
         assert "loss_mask" not in result
 
     def test_empty_text_returns_none(self, test_tokenizer):
         config = make_text_config()
-        builder = SectionedMaskBuilder()
+        builder = TextMaskBuilder()
         assert builder.build({"text": ""}, config, test_tokenizer) is None
         assert builder.build({"text": "   "}, config, test_tokenizer) is None
 
     def test_too_short_text(self, test_tokenizer):
         config = PipelineConfig(
-            input=InputConfig(sections=_TEXT_SECTIONS),
+            input=InputConfig(type="text", text_key="text"),
             preprocessing=ProcessingConfig(min_chars=100),
         )
-        builder = SectionedMaskBuilder()
+        builder = TextMaskBuilder()
         assert builder.build({"text": "short"}, config, test_tokenizer) is None
 
     def test_truncation(self, test_tokenizer):
         config = PipelineConfig(
-            input=InputConfig(sections=_TEXT_SECTIONS),
+            input=InputConfig(type="text", text_key="text"),
             preprocessing=ProcessingConfig(max_seq_len=3, min_chars=1),
         )
-        builder = SectionedMaskBuilder()
+        builder = TextMaskBuilder()
         item = {"text": "This is a very long text that should be truncated"}
         result = builder.build(item, config, test_tokenizer)
-        assert len(result["sequence"]) <= 3
+        assert len(result["ids"]) <= 3
 
 
 class TestPipeline:
@@ -410,7 +396,14 @@ class TestPipeline:
         with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
             json.dump(
                 {
-                    "special_tokens": _SPECIAL_TOKENS_CONFIG,
+                    "special_tokens": {
+                        "bos_token": "<|begin_of_sentence|>",
+                        "eos_token": "<|end_of_sentence|>",
+                        "pad_token": "<pad>",
+                        "unk_token": "<unk>",
+                        "im_start": "<|im_start|>",
+                        "im_end": "<|im_end|>",
+                    },
                     "chat_template": _CHAT_TEMPLATE,
                 },
                 f,
@@ -443,10 +436,10 @@ class TestPipeline:
             )
 
         config = PipelineConfig(
-            input=InputConfig(sections=_CHAT_SECTIONS),
+            input=InputConfig(type="chat", messages_key="messages"),
             mask={"system": "mask", "user": "mask", "assistant": "train"},
             mask_default="mask",
-            preprocessing=ProcessingConfig(max_seq_len=2048),
+            preprocessing=ProcessingConfig(max_seq_len=2048, deduplicate=True),
             output=OutputConfig(storage_format="bin", domain_key=None),
         )
 
@@ -458,16 +451,15 @@ class TestPipeline:
             tokenizer_path=tokenizer_dir,
         ).run()
 
-        meta_path = os.path.join(out_dir, "__default__", "shard_0000", "meta.json")
+        meta_path = os.path.join(out_dir, "__default__", "meta.json")
         assert os.path.exists(meta_path)
         with open(meta_path, "r") as f:
             meta = json.load(f)
         assert "sequence" in meta
         assert "loss_mask" in meta
-        assert meta["sequence"]["dtype"] == "int32"
-        assert meta["loss_mask"]["dtype"] == "int32"
 
     def test_full_text_pipeline(self, temp_dir, test_tokenizer):
+        import tempfile as tmp
 
         tokenizer_dir = os.path.join(temp_dir, "tok")
         os.makedirs(tokenizer_dir, exist_ok=True)
@@ -475,13 +467,7 @@ class TestPipeline:
         test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
         with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
             json.dump(
-                {
+                {"special_tokens": {"pad_token": "<pad>", "unk_token": "<unk>"}}, f
-                    "special_tokens": {
-                        "pad_token": "<|_pad_|>",
-                        "unk_token": "<|_unk_|>",
-                    }
-                },
-                f,
             )
 
         jsonl_path = os.path.join(temp_dir, "text.jsonl")
@@ -504,8 +490,10 @@ class TestPipeline:
             )
 
         config = PipelineConfig(
-            input=InputConfig(sections=_TEXT_SECTIONS),
+            input=InputConfig(type="text", text_key="text"),
-            preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=10),
+            preprocessing=ProcessingConfig(
+                max_seq_len=2048, min_chars=10, deduplicate=True
+            ),
             output=OutputConfig(storage_format="bin"),
         )
 
@@ -517,13 +505,12 @@ class TestPipeline:
             tokenizer_path=tokenizer_dir,
         ).run()
 
-        meta_path = os.path.join(out_dir, "__default__", "shard_0000", "meta.json")
+        meta_path = os.path.join(out_dir, "__default__", "meta.json")
         assert os.path.exists(meta_path)
         with open(meta_path, "r") as f:
             meta = json.load(f)
         assert "sequence" in meta
         assert "loss_mask" not in meta
-        assert meta["sequence"]["dtype"] == "int32"
 
     def test_full_instruction_pipeline(self, temp_dir, test_tokenizer):
         tokenizer_dir = os.path.join(temp_dir, "tok")
@@ -531,13 +518,7 @@ class TestPipeline:
         test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
         with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
             json.dump(
-                {
+                {"special_tokens": {"pad_token": "<pad>", "unk_token": "<unk>"}}, f
-                    "special_tokens": {
-                        "pad_token": "<|_pad_|>",
-                        "unk_token": "<|_unk_|>",
-                    }
-                },
-                f,
             )
 
         jsonl_path = os.path.join(temp_dir, "instruct.jsonl")
@@ -562,7 +543,9 @@ class TestPipeline:
             )
 
         config = PipelineConfig(
-            input=InputConfig(sections=_INSTRUCTION_SECTIONS),
+            input=InputConfig(
+                type="instruction", prompt_key="prompt", response_key="response"
+            ),
             mask={"prompt": "mask", "response": "train"},
             mask_default="mask",
             preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -577,66 +560,12 @@ class TestPipeline:
             tokenizer_path=tokenizer_dir,
         ).run()
 
-        meta_path = os.path.join(out_dir, "__default__", "shard_0000", "meta.json")
+        meta_path = os.path.join(out_dir, "__default__", "meta.json")
         assert os.path.exists(meta_path)
         with open(meta_path, "r") as f:
             meta = json.load(f)
         assert "sequence" in meta
         assert "loss_mask" in meta
-        assert meta["sequence"]["dtype"] == "int32"
-        assert meta["loss_mask"]["dtype"] == "int32"
-
-    def test_dtype_override(self, temp_dir, test_tokenizer):
-        tokenizer_dir = os.path.join(temp_dir, "tok")
-        os.makedirs(tokenizer_dir, exist_ok=True)
-        test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
-        with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
-            json.dump(
-                {
-                    "special_tokens": {
-                        "pad_token": "<|_pad_|>",
-                        "unk_token": "<|_unk_|>",
-                    }
-                },
-                f,
-            )
-
-        jsonl_path = os.path.join(temp_dir, "data.jsonl")
-        with open(jsonl_path, "w", encoding="utf-8") as f:
-            f.write(
-                json.dumps(
-                    {
-                        "prompt": "Q",
-                        "response": "A",
-                    }
-                )
-                + "\n"
-            )
-
-        config = PipelineConfig(
-            input=InputConfig(sections=_INSTRUCTION_SECTIONS),
-            mask={"prompt": "mask", "response": "train"},
-            mask_default="mask",
-            preprocessing=ProcessingConfig(max_seq_len=2048),
-            output=OutputConfig(
-                storage_format="bin",
-                dtype={"loss_mask": "bool"},
-            ),
-        )
-
-        out_dir = os.path.join(temp_dir, "output")
-        Pipeline(
-            config=config,
-            input_paths=[jsonl_path],
-            output_dir=out_dir,
-            tokenizer_path=tokenizer_dir,
-        ).run()
-
-        meta_path = os.path.join(out_dir, "__default__", "shard_0000", "meta.json")
-        with open(meta_path, "r") as f:
-            meta = json.load(f)
-        assert meta["sequence"]["dtype"] == "int32"
-        assert meta["loss_mask"]["dtype"] == "bool"
 
 
 class TestUtility:
@@ -646,68 +575,29 @@ class TestUtility:
         assert not filter_by_length("x" * 100, max_len=50)
         assert filter_by_length("just right", min_len=5, max_len=20)
 
-
+    def test_dedup_signature(self):
-class TestSectionedMaskBuilder:
+        a = {"key": "value", "number": 1}
-    def test_sectioned_chat(self, chat_tokenizer):
+        b = {"number": 1, "key": "value"}
-        config = PipelineConfig(
+        assert dedup_signature(a) == dedup_signature(b)
-            input=InputConfig(sections=_CHAT_SECTIONS),
+        c = {"key": "different"}
-            mask={"system": "mask", "user": "mask", "assistant": "train"},
+        assert dedup_signature(a) != dedup_signature(c)
-            mask_default="mask",
-            preprocessing=ProcessingConfig(max_seq_len=2048),
-        )
-        builder = SectionedMaskBuilder()
-        item = {
-            "messages": [
-                {"role": "user", "content": "What is 2+2?"},
-                {"role": "assistant", "content": "4"},
-            ]
-        }
-        result = builder.build(item, config, chat_tokenizer)
-        assert result is not None
-        assert len(result["sequence"]) == len(result["loss_mask"])
-        assert sum(result["loss_mask"]) > 0
-        assert 0 in result["loss_mask"]
-
-    def test_sectioned_instruction(self, test_tokenizer):
-        config = PipelineConfig(
-            input=InputConfig(sections=_INSTRUCTION_SECTIONS),
-            preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0),
-        )
-        builder = SectionedMaskBuilder()
-        item = {"prompt": "Q: Why?", "response": "A: Because."}
-        result = builder.build(item, config, test_tokenizer)
-        assert result is not None
-        mask = result["loss_mask"]
-        assert mask[0] == 0
-        assert mask[-1] == 1
-
-    def test_sectioned_text(self, test_tokenizer):
-        config = PipelineConfig(
-            input=InputConfig(sections=_TEXT_SECTIONS),
-            preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1),
-        )
-        builder = SectionedMaskBuilder()
-        item = {"text": "Hello world, this is a test."}
-        result = builder.build(item, config, test_tokenizer)
-        assert result is not None
-        assert "loss_mask" not in result
-
-    def test_sectioned_text_too_short(self, test_tokenizer):
-        config = PipelineConfig(
-            input=InputConfig(sections=_TEXT_SECTIONS),
-            preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100),
-        )
-        builder = SectionedMaskBuilder()
-        item = {"text": "short"}
-        result = builder.build(item, config, test_tokenizer)
-        assert result is None
 
 
 class TestFactoryRegistration:
     def test_registered_builders(self):
         names = MaskBuilderFactory._registry.list_names()
-        assert "sectioned" in names
+        assert "chat" in names
+        assert "instruction" in names
+        assert "text" in names
 
-    def test_create_sectioned_builder(self):
+    def test_create_chat_builder(self):
-        builder = MaskBuilderFactory.create("sectioned")
+        builder = MaskBuilderFactory.create("chat")
-        assert isinstance(builder, SectionedMaskBuilder)
+        assert isinstance(builder, ChatMaskBuilder)
+
+    def test_create_instruction_builder(self):
+        builder = MaskBuilderFactory.create("instruction")
+        assert isinstance(builder, InstructionMaskBuilder)
+
+    def test_create_text_builder(self):
+        builder = MaskBuilderFactory.create("text")
+        assert isinstance(builder, TextMaskBuilder)