Data

This section covers the pipeline from parquet files containing CIF text to the batches consumed during training. It includes the tokenizer that converts CIF strings to token IDs, the dataset loader that tokenizes and filters the data, and the collator that assembles fixed-length training batches.

The preprocessing CLIs used to generate the parquet files are documented on the CLI entry points page. The models that consume these batches are documented on the Models page.

Tokenizer

The CIF tokenizer provides the common text interface used by all model families.

CustomCIFTokenizer

Bases: PreTrainedTokenizer

Hugging Face-compatible tokenizer for CIF text.

Tokenisation follows CIF structure rather than subwords, keeping element symbols, numbers, and CIF keywords as whole tokens. Space-group tokens receive an internal _sg suffix during tokenisation and the suffix is removed during decoding. from_pretrained and save_pretrained use vocabulary.json, spacegroups.txt, and tokenizer_config.json.

Source code in _tokenizer.py

class CustomCIFTokenizer(PreTrainedTokenizer):
    """Hugging Face-compatible tokenizer for CIF text.

    Tokenisation follows CIF structure rather than subwords, keeping element symbols, numbers, and
    CIF keywords as whole tokens. Space-group tokens receive an internal `_sg` suffix during
    tokenisation and the suffix is removed during decoding. `from_pretrained` and `save_pretrained`
    use `vocabulary.json`, `spacegroups.txt`, and `tokenizer_config.json`.
    """
    def __init__(
        self,
        vocab_file: str,
        spacegroups_file: str | None,
        unk_token: str="<unk>",
        pad_token: str | None="<pad>",
        bos_token: str="<bos>",
        eos_token: str="<eos>",
        var_open_token: str="[",
        var_close_token: str="]",
        prop_token: str="<prop>",
        max_vocab_size: int | None=None,
        **kwargs
    ) -> None:
        with open(vocab_file, "r") as f:
            self.token_to_id = json.load(f)

        # New tokens are appended, so truncation preserves older vocabularies.
        if max_vocab_size is not None:
            self.token_to_id = {t: i for t, i in self.token_to_id.items() if i < max_vocab_size}

        # Invert to get ID -> token
        self.id_to_token = {v: k for k, v in self.token_to_id.items()}

        # Keep track of all known tokens
        self._tokens = list(self.token_to_id.keys())

        # For convenience, load the raw spacegroups
        with open(spacegroups_file, "r") as f:
            self.space_groups = [sg.strip() for sg in f.readlines()]

        # A space group whose token was cut (P6 for a 377-token checkpoint) splits into plain tokens, as it used to
        self.space_groups = [sg for sg in self.space_groups if sg + "_sg" in self.token_to_id]

        # Sort tokens by length for regex matching
        self._escaped_tokens = sorted(
            [re.escape(t) for t in self._tokens],
            key=len,
            reverse=True
        )

        # re.escape backslashes any regex metacharacter in a token, so a token containing
        # "\\n" or "." matches as literal text instead of as a newline or a wildcard.
        # Longest first matters: with "a" ahead of "ab" in the alternation, the regex would
        # match "a" and leave a stray "b" to be tokenised separately.

        # Define main special tokens
        self.unk_token = unk_token
        self.pad_token = pad_token
        self.bos_token = bos_token
        self.eos_token = eos_token
        self.var_open_token = var_open_token
        self.var_close_token = var_close_token
        self.prop_token = prop_token

        # Make sure they're all in the vocabulary
        for special_tk in [self.unk_token, 
                           self.pad_token, 
                           self.bos_token, 
                           self.eos_token, 
                           self.var_open_token,
                           self.var_close_token,
                           self.prop_token]:
            if special_tk not in self.token_to_id:
                new_id = len(self.token_to_id)
                self.token_to_id[special_tk] = new_id
                self.id_to_token[new_id] = special_tk
                self._tokens.append(special_tk)

        # Validate id_to_token mapping
        self.validate_id_to_token()

        super().__init__(
            unk_token=self.unk_token,
            pad_token=self.pad_token,
            bos_token=self.bos_token,
            eos_token=self.eos_token,
            var_open_token=self.var_open_token,
            var_close_token=self.var_close_token,
            prop_token=self.prop_token,
            **kwargs
        )

    @property
    def vocab_size(self) -> int:
        """Return the size of the base vocabulary (without added special tokens).
        """
        return len(self.token_to_id)

    def _tokenize(self, text: str) -> list[str]:
        """Custom tokenization logic for CIF data.
        """
        # Disambiguate space groups in the text:
        spacegroups_pattern = "|".join(self.space_groups)

        # Disambiguate bracketed space groups:
        text = re.sub(
            fr'(_symmetry_space_group_name_H-M\s*\[)({spacegroups_pattern})(\])',
            r'\1\2_sg\3',
            text
        )


        # Add "_sg" to the sample's space-group name.

        # Build the tokenization pattern:
        token_pattern = "|".join(self._escaped_tokens)
        # all tokens from sample are escaped and joined by "|"
        # allows to match any of predefined tokens

        # Define the full pattern to match tokens, words, and punctuation
        full_pattern = f"({token_pattern}|\\w+|[\\.,;!?\\[\\]])"
        # can match any of the predefined tokens, words, or punctuation

        # Clean up spaces
        text = re.sub(r"[ \t]+", " ", text)
        # replace multiple spaces or tabs with a single space, gives uniform

        # Extract tokens
        tokens = re.findall(full_pattern, text)
        # find all matches of the pattern in the text

        # Replace unknown tokens
        output_tokens = [
            t if t in self._tokens else self.unk_token
            for t in tokens
        ]
        # if the token is in the predefined tokens, keep it, otherwise replace with unk_token
        return output_tokens

    def _convert_token_to_id(self, token: str) -> int:
        """Convert a token to its corresponding ID.
        """
        return self.token_to_id.get(token, self.token_to_id.get(self.unk_token))

    def _convert_id_to_token(self, index: int) -> str:
        """Convert an ID to its corresponding token.
        """
        #print the id_to_token dictionary
        if hasattr(index, "item"):
            index = int(index.item())  # Convert tensor to int
        if index not in self.id_to_token:
            print(f"Warning: Token ID {index} not found in id_to_token. Returning <unk>.")
        return self.id_to_token.get(index, self.unk_token)

    def validate_id_to_token(self) -> None:
        """Validate that id_to_token covers all token IDs up to vocab_size.
        """
        missing_ids = [i for i in range(len(self.token_to_id)) if i not in self.id_to_token]
        if missing_ids:
            print(f"Warning: Missing IDs in id_to_token: {missing_ids}")
        else:
            print("Tokenizer validation passed: token vocabulary is consistent.")

    def convert_tokens_to_string(self, tokens: list[str]) -> str:
        """Convert a list of tokens to a string.
        """
        return "".join(tokens)

    def build_inputs_with_special_tokens(self, token_ids_0: list[int], token_ids_1: list[int] | None=None) -> list[int]:
        """Add special tokens to a sequence or a pair of sequences. For GPT-2-style models, no additional
special tokens are used. either returns token_ids_0 or token_ids_0 + token_ids_1
        """
        if token_ids_1 is None:
            return token_ids_0
        return token_ids_0 + token_ids_1

    def get_vocab(self) -> dict[str, int]:
        """Return the tokenizer's vocabulary (including base vocabulary).
        """
        return dict(self.token_to_id)

    def get_added_vocab(self) -> dict[str, int]:
        """Return the additional vocabulary tokens added after the initial vocab.
        """
        return {
            tok: idx
            for tok, idx in self._added_tokens_encoder.items()
        }

    def decode(self, token_ids: list[int] | torch.Tensor, skip_special_tokens: bool=False, **kwargs) -> str:
        """Decode token ids into CIF text.

        Internal `_sg` suffixes are removed from space-group tokens so the decoded text uses the
        external CIF spelling.
        """
        tokens = [self._convert_id_to_token(idx) for idx in token_ids]
        if skip_special_tokens:
            tokens = [tok for tok in tokens if tok not in [self.bos_token, 
                                                           self.eos_token, 
                                                           self.pad_token, 
                                                           self.var_open_token, 
                                                           self.var_close_token,
                                                           self.prop_token]]
        # Remove '_sg' suffix from space group tokens
        tokens = [re.sub(r'_sg$', '', tok) for tok in tokens]

        result = self.convert_tokens_to_string(tokens)
        return result

    @classmethod
    def from_pretrained(cls, pretrained_dir: str, **kwargs) -> "CustomCIFTokenizer":
        """Load a tokenizer from a directory written by `save_pretrained`.

        The directory must contain `vocabulary.json`, `spacegroups.txt`, and `tokenizer_config.json`. Use `max_vocab_size` to match an older checkpoint's vocabulary.
        """
        pretrained_dir = resolve_tokenizer_dir(pretrained_dir)
        vocab_file = os.path.join(pretrained_dir, "vocabulary.json")
        spacegroups_file = os.path.join(pretrained_dir, "spacegroups.txt")
        tokenizer_config_file = os.path.join(pretrained_dir, "tokenizer_config.json")

        with open(tokenizer_config_file, "r") as f:
            config = json.load(f)

        return cls(
            vocab_file=vocab_file,
            spacegroups_file=spacegroups_file,
            unk_token=config.get("unk_token", "<unk>"),
            **kwargs
        )

    def save_vocabulary(self, save_directory: str, filename_prefix: str | None=None) -> tuple[str, ...]:
        """Save the base vocabulary (token->ID) to a file. Hugging Face's `save_pretrained` will call this.
Returns the path(s) of the saved vocab file(s).
        """
        if not os.path.isdir(save_directory):
            os.makedirs(save_directory)

        vocab_file_name = (filename_prefix + "-" if filename_prefix else "") + "vocabulary.json"
        vocab_file = os.path.join(save_directory, vocab_file_name)
        with open(vocab_file, "w") as f:
            json.dump(self.token_to_id, f)

        return (vocab_file,)

    def save_pretrained(self, save_directory: str, **kwargs) -> None:
        """Save the tokenizer vocabulary, space groups, and configuration.

        The files are written in the layout expected by `from_pretrained`.
        """
        super().save_pretrained(save_directory, **kwargs)

        # Also save spacegroups
        spacegroups_file = os.path.join(save_directory, "spacegroups.txt")
        with open(spacegroups_file, "w") as f:
            f.write("\n".join(self.space_groups))

        # Save tokenizer configuration
        tokenizer_config = {
            "unk_token": self.unk_token,
            "vocab_size": self.vocab_size,
            "tokenizer_class": self.__class__.__name__,
        }
        tokenizer_config_file = os.path.join(save_directory, "tokenizer_config.json")
        with open(tokenizer_config_file, "w") as f:
            json.dump(tokenizer_config, f)

    def add_custom_tokens(self, tokens: list[str]) -> None:
        """Add tokens to the vocabulary and refresh the tokenisation pattern.

        New tokens are appended so existing token ids retain their meanings and remain compatible
        with models trained against the previous vocabulary.
        """
        for token in tokens:
            if token not in self.token_to_id:
                new_id = len(self.token_to_id)
                self.token_to_id[token] = new_id
                self.id_to_token[new_id] = token
                self._tokens.append(token)  # Add to the list of tokens
                print(f"Added token '{token}' with ID {new_id}")

        # Update escaped tokens for regex matching
        self._escaped_tokens = sorted(
            [re.escape(t) for t in self._tokens],
            key=len,
            reverse=True
        )

    def remove_custom_tokens(self, tokens: list[str]) -> None:
        """Remove tokens from the vocabulary.

        This is safe only before training. Removing a token shifts subsequent token ids, which makes
        models trained against the previous vocabulary incompatible with the tokenizer.
        """
        removed_any = False
        for token in tokens:
            if token in self.token_to_id:
                # Don't remove special tokens
                if token in [self.unk_token, self.pad_token, self.bos_token, 
                           self.eos_token, self.var_open_token, self.var_close_token, self.prop_token]:
                    print(f"Skipping special token '{token}' - cannot remove")
                    continue

                old_id = self.token_to_id.pop(token)
                del self.id_to_token[old_id]
                self._tokens.remove(token)
                removed_any = True
                print(f"Removed token '{token}' (was ID {old_id})")

        if removed_any:
            # Update escaped tokens only. Downstream code uses the original IDs.
            self._escaped_tokens = sorted(
                [re.escape(t) for t in self._tokens],
                key=len,
                reverse=True
            )
            print("WARNING: Token removal creates gaps in ID space. Model compatibility may be affected.")

    def validate_tokenizer_state(self) -> None:
        """Validate that the tokenizer's internal state is consistent.
        """
        # Check token_to_id and id_to_token mappings
        for token, idx in self.token_to_id.items():
            if self.id_to_token.get(idx) != token:
                print(f"Warning: Inconsistent mapping for token '{token}' (ID {idx})")

        # Check _tokens list
        missing_tokens = [token for token in self.token_to_id.keys() if token not in self._tokens]
        if missing_tokens:
            print(f"Warning: Missing tokens in _tokens: {missing_tokens}")

        # Check _escaped_tokens list
        missing_escaped = [re.escape(token) for token in self.token_to_id.keys() if re.escape(token) not in self._escaped_tokens]
        if missing_escaped:
            print(f"Warning: Missing escaped tokens in _escaped_tokens: {missing_escaped}")

        print("Tokenizer state validation complete.") 

Dataset loading

load_data tokenizes a Hugging Face dataset and returns the tokenized dataset together with the collator required by the Hugging Face Trainer. For conditional runs, condition_columns specifies the dataset columns containing the conditioning values. These values are preserved through tokenization and passed through to the training batch.

Filtering is optional. By default, CIFs longer than the context length are truncated (but i typically prefer to do this in preprocessing), and CIFs containing unknown tokens are retained. Set remove_CIFs_above_context or remove_CIFs_with_unk to discard these samples instead.

load_data

load_data(tokenizer, dataset, context_length, mode='unconditional', condition_columns=None, remove_CIFs_above_context=False, remove_CIFs_with_unk=False, show_token_stats=False, validate_conditions=False)

Tokenize a CIF dataset and build the collator used by the Hugging Face Trainer.

Conditional datasets must provide the requested condition columns. CIFs exceeding the context length or containing unknown tokens can be removed instead of being truncated or retained.

Parameters:

Name	Type	Description	Default
tokenizer	CustomCIFTokenizer	CIF tokenizer used to encode structures and provide the pad token id.	required
dataset	object	Hugging Face dataset containing CIF text and, for conditional training, the required condition columns.	required
context_length	int	Number of tokens in each training sequence.	required
mode	str	Whether training is unconditional or conditional.	'unconditional'
condition_columns	list[str] | None	Dataset columns containing the conditioning values for conditional training.	None
remove_CIFs_above_context	bool	Whether to drop CIFs longer than the context length.	False
remove_CIFs_with_unk	bool	Whether to drop CIFs containing unknown tokens.	False
show_token_stats	bool	Whether to print token-length statistics.	False
validate_conditions	bool	Whether to validate conditioning values before training.	False

Returns:

Type	Description
tuple	The tokenized dataset and data collator ready for the Hugging Face Trainer.

Source code in _dataloader.py

def load_data(
    tokenizer: "CustomCIFTokenizer", 
    dataset: object, 
    context_length: int, 
    mode: str="unconditional",
    condition_columns: list[str] | None=None,
    remove_CIFs_above_context: bool=False, 
    remove_CIFs_with_unk: bool=False,
    show_token_stats: bool=False,
    validate_conditions: bool=False
) -> tuple:
    """Tokenize a CIF dataset and build the collator used by the Hugging Face Trainer.

    Conditional datasets must provide the requested condition columns. CIFs exceeding the context
    length or containing unknown tokens can be removed instead of being truncated or retained.

    Args:
        tokenizer: CIF tokenizer used to encode structures and provide the pad token id.
        dataset: Hugging Face dataset containing CIF text and, for conditional training, the
                 required condition columns.
        context_length: Number of tokens in each training sequence.
        mode: Whether training is unconditional or conditional.
        condition_columns: Dataset columns containing the conditioning values for conditional
                           training.
        remove_CIFs_above_context: Whether to drop CIFs longer than the context length.
        remove_CIFs_with_unk: Whether to drop CIFs containing unknown tokens.
        show_token_stats: Whether to print token-length statistics.
        validate_conditions: Whether to validate conditioning values before training.

    Returns:
        The tokenized dataset and data collator ready for the Hugging Face Trainer.
    """

    # Validate inputs
    if mode == "conditional" and condition_columns is None:
        raise ValueError(f"condition_columns must be provided for mode='{mode}'")

    # Parse condition columns for conditional mode
    parsed_condition_columns = None
    if mode == "conditional":
        try:
            parsed_condition_columns = ast.literal_eval(str(condition_columns))
        except Exception as e:
            raise ValueError(f"Error parsing condition_columns: {condition_columns}. {e}")

    # Create tokenizer function based on mode
    if mode == "unconditional":
        def tokenize_fn(examples):
            return tokenize_function(examples, tokenizer, context_length, mode="unconditional")
    elif mode == "conditional":
        def tokenize_fn(examples):
            return tokenize_function(examples, tokenizer, condition_columns, mode="conditional")
    else:
        raise ValueError(f"Unknown mode: {mode}")

    # Handle dataset indexing for conditional mode
    if mode == "conditional" and validate_conditions:
        dataset_with_idx = dataset.map(
            lambda ex, idx: {"__raw_idx": idx},
            with_indices=True,
            batched=False
        )
    else:
        dataset_with_idx = dataset

    # Determine columns to keep based on mode
    if mode == "conditional":
        columns_to_keep = [
            "input_ids",
            "attention_mask",
            "fixed_mask",
            "special_tokens_mask",
            "condition_values",
        ]
        if validate_conditions:
            columns_to_keep.append("__raw_idx")
    else:  # unconditional
        columns_to_keep = [
            "input_ids",
            "token_type_ids",
            "attention_mask",
            "fixed_mask",
            "special_tokens_mask",
        ]

    # Calculate columns to remove (reduce memory usage)
    original_columns = dataset_with_idx["train"].column_names
    if mode == "conditional":
        columns_to_remove = [
            col for col in original_columns 
            if col not in columns_to_keep and col not in parsed_condition_columns and col != "CIF"
        ]
    else:
        columns_to_remove = [col for col in original_columns if col not in columns_to_keep]

    # Tokenize dataset
    tokenized_dataset = dataset_with_idx.map(
        tokenize_fn,
        batched=True,
        remove_columns=columns_to_remove,
        num_proc=NUM_PROC_TOK
    )


    # Few filters/validations/stats
    if remove_CIFs_above_context:
        print("\n******* Tokenization Filters Stats *******")
        print(f"Removing CIFs with token length exceeding {context_length}")
        tokenized_dataset = filter_long_CIFs(tokenized_dataset, context_length)

    if remove_CIFs_with_unk:
        print("Removing CIFs with unknown tokens")
        tokenized_dataset = filter_CIFs_with_unk(tokenized_dataset, tokenizer)

    if "fixed_mask" not in tokenized_dataset["train"].features:
        raise ValueError("Fixed mask not present in tokenized dataset")

    if show_token_stats:
        print("\n******* Token Length Stats *******")
        get_token_length_stats(tokenized_dataset)

    if mode == "conditional" and validate_conditions:
        validate_condition_values(tokenized_dataset, dataset_with_idx, parsed_condition_columns)

    data_collator = CustomCIFDataCollator(tokenizer, context_length)
    return tokenized_dataset, data_collator

Custom Dataloader for Conditioning

CustomCIFDataCollator constructs the batches consumed during training. Changes to the training pipeline may require changes here.

Every sequence produced by the collator has exactly context_length tokens. If a CIF exceeds the context window, it is truncated from the beginning. Shorter CIFs are then packed with tokens from other CIFs, selected round-robin from the remainder of the batch, until the window is filled. This packing keeps the context window occupied by CIF tokens rather than padding, but a single sequence can therefore contain parts of multiple structures.

The collator determines whether a batch is conditional from the features themselves by checking for condition_values. The same collator therefore handles both conditional and unconditional training without an additional mode flag. For a packed sequence, the associated condition values come from the CIF that provided the start of the sequence rather than from every CIF contributing tokens to it.

Each batch is returned as a dictionary containing [B, context_length] tensors for input_ids, labels, fixed_mask, attention_mask, and special_tokens_mask. Padding positions in labels are set to -100. Conditional batches also contain condition_values.

Scalar conditions are stacked into tensors of shape [B, n_conditions]. Continuous XRD conditions can contain a different number of rows for each sample, so they are front-padded with -100.0 to match the longest sample in the batch before stacking.

CustomCIFDataCollator

Collate tokenized CIF rows into fixed-length training batches.

Every sequence in a batch is exactly context_length tokens. A CIF longer than that is sliced from the beginning, a shorter one is packed with further CIFs taken round-robin from the rest of the batch until the window is full. Packing keeps the context window full of real tokens instead of padding.

Conditional mode is detected from the features themselves, by whether condition_values is present, so the same collator serves both training modes.

Source code in _dataloader.py

class CustomCIFDataCollator:
    """Collate tokenized CIF rows into fixed-length training batches.

    Every sequence in a batch is exactly `context_length` tokens. A CIF longer than that is sliced
    from the beginning, a shorter one is packed with further CIFs taken round-robin from the rest of
    the batch until the window is full. Packing keeps the context window full of real tokens instead
    of padding.

    Conditional mode is detected from the features themselves, by whether `condition_values` is
    present, so the same collator serves both training modes.
    """
    def __init__(self, tokenizer: "CustomCIFTokenizer", context_length: int, pack: bool = True) -> None:
        self.tokenizer = tokenizer
        self.context_length = context_length
        self.pack = pack  # False keeps CIFs separate for evaluation.

    def _pad_unpacked(self, features: list[dict]) -> dict[str, torch.Tensor]:
        """One CIF per row, truncated at `context_length`, with masked right-padding to the longest row."""
        rows = [list(feature["input_ids"])[:self.context_length] for feature in features]
        width = max(len(row) for row in rows)

        def pad(values, fill):
            values = list(values)[:self.context_length]
            return values + [fill] * (width - len(values))

        input_ids = torch.tensor([pad(row, self.tokenizer.pad_token_id) for row in rows], dtype=torch.long)
        attention_mask = torch.tensor([pad([1] * len(row), 0) for row in rows], dtype=torch.long)

        batch = {
            "input_ids": input_ids,
            "labels": input_ids.masked_fill(attention_mask == 0, -100),
            "fixed_mask": torch.tensor([pad(feature["fixed_mask"], 0) for feature in features], dtype=torch.long),
            "attention_mask": attention_mask,
            "special_tokens_mask": torch.zeros_like(input_ids),
        }
        if "condition_values" in features[0]:
            batch["condition_values"] = _pack_condition_values([feature["condition_values"] for feature in features])
        return batch

    def __call__(self, features: list[dict]) -> dict[str, torch.Tensor]:
        """Pack tokenised features into a fixed-length batch.

        Conditional mode is determined from the first feature containing
        ``condition_values``. A packed sequence without ``context_length`` raises
        ``ValueError`` because the resulting batch would be ragged.

        Args:
            features: Tokenised rows containing ``input_ids``, ``fixed_mask``, and
                optionally ``attention_mask``, ``special_tokens_mask``, and
                ``condition_values``.

        Returns:
            A dictionary of ``[B, context_length]`` tensors containing ``input_ids``,
            ``labels``, ``fixed_mask``, ``attention_mask``, and
            ``special_tokens_mask``. Conditional batches also contain
            ``condition_values``. Padding positions in ``labels`` are set to ``-100``.
        """

        if not self.pack:
            return self._pad_unpacked(features)

        # Auto-detect conditional mode based on presence of condition_values
        is_conditional = "condition_values" in features[0]

        # Prepare batch containers
        batch_input_ids = []
        batch_fixed_mask = []
        batch_attention_mask = []
        batch_special_tokens_mask = []
        batch_condition_values = [] if is_conditional else None

        def slice_from_beginning(input_ids, fixed_mask, special_tokens_mask, attention_mask):
            if len(input_ids) > self.context_length:
                input_ids = input_ids[:self.context_length]
                fixed_mask = fixed_mask[:self.context_length]
                if special_tokens_mask is not None:
                    special_tokens_mask = special_tokens_mask[:self.context_length]
                attention_mask = attention_mask[:self.context_length]
            return input_ids, fixed_mask, special_tokens_mask, attention_mask

        # Convert to Python lists for indexing
        for feature in features:
            feature["input_ids"] = list(feature["input_ids"])
            feature["fixed_mask"] = list(feature["fixed_mask"])
            feature["attention_mask"] = list(feature.get("attention_mask", [1]*len(feature["input_ids"])))
            feature["special_tokens_mask"] = list(feature.get("special_tokens_mask", [0] * len(feature["input_ids"])))

        # Pack sequences
        for i in range(len(features)):
            if is_conditional:
                current_condition_values = features[i]["condition_values"]

            input_ids_i = features[i]["input_ids"]
            fixed_mask_i = features[i]["fixed_mask"]
            special_mask_i = features[i]["special_tokens_mask"]
            attention_mask_i = features[i]["attention_mask"]

            # Single long CIF - slice from beginning
            if len(input_ids_i) > self.context_length:
                packed_input_ids, packed_fixed_mask, packed_special_tokens_mask, packed_attention_mask = (
                    slice_from_beginning(input_ids_i, fixed_mask_i, special_mask_i, attention_mask_i)
                )
            else:
                # Pack multiple CIFs
                packed_input_ids = []
                packed_fixed_mask = []
                packed_special_tokens_mask = []
                packed_attention_mask = []

                current_idx = i
                while len(packed_input_ids) < self.context_length:
                    block_input_ids = features[current_idx]["input_ids"]
                    block_fixed_mask = features[current_idx]["fixed_mask"]
                    block_special_mask = features[current_idx]["special_tokens_mask"]
                    block_attention = features[current_idx]["attention_mask"]

                    # Trim block if it would exceed context_length
                    new_total_length = len(packed_input_ids) + len(block_input_ids)
                    if new_total_length > self.context_length:
                        space_left = self.context_length - len(packed_input_ids)
                        block_input_ids = block_input_ids[:space_left]
                        block_fixed_mask = block_fixed_mask[:space_left]
                        if block_special_mask is not None:
                            block_special_mask = block_special_mask[:space_left]
                        block_attention = block_attention[:space_left]

                    # Fill missing special_tokens_mask with zeros
                    if block_special_mask is None:
                        block_special_mask = [0] * len(block_input_ids)

                    # Extend packed sequences
                    packed_input_ids.extend(block_input_ids)
                    packed_fixed_mask.extend(block_fixed_mask)
                    packed_special_tokens_mask.extend(block_special_mask)
                    packed_attention_mask.extend(block_attention)

                    # Break if we've reached context_length
                    if len(packed_input_ids) >= self.context_length:
                        packed_input_ids = packed_input_ids[:self.context_length]
                        packed_fixed_mask = packed_fixed_mask[:self.context_length]
                        packed_special_tokens_mask = packed_special_tokens_mask[:self.context_length]
                        packed_attention_mask = packed_attention_mask[:self.context_length]
                        break

                    # Move to next feature in round-robin
                    current_idx = (current_idx + 1) % len(features)

                # Add special_tokens_mask when missing
                if len(packed_special_tokens_mask) == 0:
                    packed_special_tokens_mask = [0] * len(packed_input_ids)

            # validation
            if len(packed_input_ids) != self.context_length:
                raise ValueError("Packed sequence length mismatch")

            # Add to batch
            batch_input_ids.append(packed_input_ids)
            batch_fixed_mask.append(packed_fixed_mask)
            batch_attention_mask.append(packed_attention_mask)
            batch_special_tokens_mask.append(packed_special_tokens_mask)

            if is_conditional:
                batch_condition_values.append(current_condition_values)

        # Convert to tensors
        batch_input_ids = torch.tensor(batch_input_ids, dtype=torch.long)
        batch_fixed_mask = torch.tensor(batch_fixed_mask, dtype=torch.long)
        batch_attention_mask = torch.tensor(batch_attention_mask, dtype=torch.long)
        batch_special_tokens_mask = torch.tensor(batch_special_tokens_mask, dtype=torch.long)

        # Labels
        labels = batch_input_ids.clone()
        labels[labels == self.tokenizer.pad_token_id] = -100

        # Build output batch
        batch = {
            "input_ids": batch_input_ids,
            "labels": labels,
            "fixed_mask": batch_fixed_mask,
            "attention_mask": batch_attention_mask,
            "special_tokens_mask": batch_special_tokens_mask
        }

        # Add condition values for conditional mode
        if is_conditional:
            batch["condition_values"] = _pack_condition_values(batch_condition_values)

        return batch