mostlyai-engine 1.0.0__py3-none-any.whl

mostlyai/engine/_language/generation.py

@@ -0,0 +1,309 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+
+import json_repair
+import logging
+import time
+from pathlib import Path
+
+import pandas as pd
+import torch
+from huggingface_hub import constants as hf_constants
+from transformers import (
+    PreTrainedTokenizerBase,
+)
+
+from mostlyai.engine._common import (
+    persist_data_part,
+    FixedSizeSampleBuffer,
+    STRING,
+    ProgressCallback,
+    ProgressCallbackWrapper,
+)
+from mostlyai.engine._language.common import estimate_max_tokens, MAX_LENGTH
+from mostlyai.engine._language.encoding import encode_df
+from mostlyai.engine._workspace import ensure_workspace_dir, Workspace, reset_dir
+from mostlyai.engine._language.formatron_utils import (
+    get_formatter_builders,
+    prepare_seed_for_formatron,
+    get_vocab_processors,
+)
+
+INVALID_VALUE = "_INVALID_"  # when JSON parsing fails, the values of target columns will be set to this
+DUMMY_CONTEXT_KEY = "__dummy_context_key"
+_LOG = logging.getLogger(__name__)
+
+
+def decode_buffered_samples(
+    buffer: FixedSizeSampleBuffer,
+    tokenizer: PreTrainedTokenizerBase,
+    tgt_text_columns: list[str],
+    tgt_context_key: str,
+    max_new_tokens: int,
+):
+    t0 = time.time()
+
+    def parse_json(x, columns: list[str]):
+        try:
+            parsed_x = json_repair.loads(x)
+            if not isinstance(parsed_x, dict):
+                raise ValueError("parsed_x has to be a dictionary")
+        except (json.decoder.JSONDecodeError, ValueError):
+            parsed_x = {}
+        return [parsed_x.get(c, INVALID_VALUE) for c in columns]
+
+    ctx_keys = []
+    tgt_seed = []
+    output_texts = []
+    num_samples_max_length_limit = 0
+    for outputs_ids, keys_df, seed_df in buffer.buffer:
+        try:
+            num_tokens_by_row = [sum(token != tokenizer.eos_token_id for token in row) for row in outputs_ids]
+            num_samples_max_length_limit += sum(1 for tokens in num_tokens_by_row if tokens >= max_new_tokens)
+        except AttributeError:
+            num_samples_max_length_limit = float("-inf")
+        outputs_text = tokenizer.batch_decode(outputs_ids, skip_special_tokens=True)
+        output_texts.extend(outputs_text)
+        ctx_keys.append(keys_df)
+        tgt_seed.append(seed_df)
+    _LOG.info(f"{num_samples_max_length_limit=}")
+    ctx_keys = pd.concat(ctx_keys, axis=0).reset_index(drop=True).rename(tgt_context_key)
+    tgt_seed = pd.concat(tgt_seed, axis=0).reset_index(drop=True)
+    # The model works with un-prefixed column names, but we need to recover prefixed column names for the final output
+    tgt_data = pd.DataFrame(
+        [parse_json(text, tgt_text_columns) for text in output_texts],
+        columns=tgt_text_columns,
+        index=ctx_keys.index,
+        dtype="string",
+    )
+    # make sure invalid/incomplete unicode chars are replaced with the replacement char � (U+FFFD)
+    tgt_data = tgt_data.map(
+        lambda x: x.encode("utf-8", errors="surrogatepass").decode("utf-8", errors="replace") if not pd.isna(x) else x
+    )
+    # overwrite generated columns with the seeded values
+    tgt_data.update(tgt_seed)
+    # ensure STRING type
+    tgt_data = tgt_data.astype(STRING)
+    # prepend the context keys to the data (if not dummy context)
+    if ctx_keys.name != DUMMY_CONTEXT_KEY:
+        tgt_data = pd.concat([ctx_keys, tgt_data], axis=1)
+    invalid_percentage = ((tgt_data[tgt_text_columns] == INVALID_VALUE).sum() / len(tgt_data) * 100.0).map(
+        "{:.2f}%".format
+    )
+    _LOG.info(f"percentage of invalid values: {invalid_percentage.to_dict()}")
+    _LOG.info(f"decoded {tgt_data.shape} from {len(buffer.buffer)} batches in {time.time() - t0:.2f}s")
+    return tgt_data
+
+
+def generate(
+    *,
+    ctx_data: pd.DataFrame | None = None,
+    seed_data: pd.DataFrame | None = None,
+    sample_size: int | None = None,
+    batch_size: int | None = None,
+    sampling_temperature: float = 1.0,
+    sampling_top_p: float = 1.0,
+    device: torch.device | str | None = None,
+    workspace_dir: str | Path = "engine-ws",
+    update_progress: ProgressCallback | None = None,
+):
+    _LOG.info("GENERATE_LANGUAGE started")
+    t0_ = time.time()
+    with ProgressCallbackWrapper(update_progress) as progress:
+        device = (
+            torch.device(device)
+            if device is not None
+            else (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
+        )
+        _LOG.info(f"{device=}")
+        _LOG.info(f"{sampling_temperature=}, {sampling_top_p=}")
+
+        workspace_dir = ensure_workspace_dir(workspace_dir)
+        workspace = Workspace(workspace_dir)
+        output_path = workspace.generated_data_path
+        reset_dir(output_path)
+        tgt_stats = workspace.tgt_stats.read()
+        tgt_text_columns = list(tgt_stats["columns"].keys())
+        tgt_context_key = tgt_stats["keys"].get("context_key")
+        has_context = workspace.ctx_stats.path.exists()
+
+        # resolve potential conflict between sample_seed and sample_size
+        if seed_data is not None:
+            assert sample_size is None, "either sample_seed or sample_size can be provided, not both"
+            sample_size = len(seed_data)
+
+        if has_context:
+            ctx_stats = workspace.ctx_stats.read()
+            ctx_columns = list(ctx_stats["columns"].keys())
+            ctx_primary_key = ctx_stats["keys"].get("primary_key")
+
+            # ensure ctx_data exists
+            if ctx_data is None:
+                if workspace.ctx_data_path.exists():
+                    # attempt to re-use context from training, if no new context provided
+                    ctx_data = pd.read_parquet(workspace.ctx_data_path)
+                else:
+                    # build dummy context; fallback to using training data size as sample_size
+                    trn_sample_size = tgt_stats["no_of_training_records"] + tgt_stats["no_of_validation_records"]
+                    ctx_data = pd.DataFrame({ctx_primary_key: list(range(sample_size or trn_sample_size))})
+            _LOG.info(f"{ctx_data.shape=}")
+            ctx_data = ctx_data.reset_index(drop=True)
+            ctx_data_len = len(ctx_data)
+            if sample_size is not None and sample_size < ctx_data_len:
+                # take first `sample_size` rows of context
+                ctx_data = ctx_data.head(sample_size)
+                _LOG.info(f"dropped {ctx_data_len - len(ctx_data)} rows from context data")
+            _LOG.info(f"{ctx_data.shape=}")
+
+            # update sample_size based on fetched context data
+            sample_size = len(ctx_data)
+            _LOG.info(f"{sample_size=}")
+        else:
+            # create on-the-fly context
+            if sample_size is None:
+                trn_sample_size = tgt_stats["no_of_training_records"] + tgt_stats["no_of_validation_records"]
+                sample_size = trn_sample_size if sample_size is None else sample_size
+            ctx_columns = []
+            ctx_primary_key = tgt_context_key = DUMMY_CONTEXT_KEY
+            ctx_data = pd.DataFrame({ctx_primary_key: range(sample_size)})
+
+        # ensure sample_seed exists; ensure valid columns
+        if seed_data is None:
+            # build dummy seed
+            seed_data = pd.DataFrame(index=list(range(sample_size)))
+        seed_data = seed_data[[c for c in tgt_text_columns if c in seed_data.columns]]
+        _LOG.info(f"{seed_data.shape=}")
+
+        # sanity check: at this point sample seed and context data should have the same number of rows
+        assert len(seed_data) == len(ctx_data)
+
+        # early exit in case generation context is empty
+        if sample_size == 0:
+            _LOG.info("terminating generation early as no context data provided")
+            empty_out_df = pd.DataFrame(columns=[tgt_context_key] + tgt_text_columns, dtype="string")
+            persist_data_part(empty_out_df, output_path, f"{0:06}.{0:06}")
+            return
+
+        # encode context data
+        encoded_ctx_data = encode_df(ctx_df=ctx_data, ctx_columns=ctx_columns)
+
+        # estimate max new tokens based on char length of original data; consider JSON overhead
+        max_new_tokens = estimate_max_tokens(tgt_stats)
+        _LOG.info(f"{max_new_tokens=}")
+
+        t0 = time.time()
+        hf_constants.HF_HUB_OFFLINE = (
+            False  # needed for gated hf models that are not sharded, otherwise GatedRepoError in vLLM
+        )
+        # set the default env var so that we don't pass it explicitly to vLLM
+        os.environ["HF_TOKEN"] = os.getenv("MOSTLY_HUGGING_FACE_TOKEN", "")
+
+        is_peft_adapter = (workspace.model_path / "adapter_config.json").exists()
+        if is_peft_adapter and device.type == "cuda":
+            from mostlyai.engine._language.engine.vllm_engine import VLLMEngine
+
+            engine = VLLMEngine(workspace.model_path, device, max_new_tokens, MAX_LENGTH)
+        else:
+            from mostlyai.engine._language.engine.hf_engine import HuggingFaceEngine
+
+            engine = HuggingFaceEngine(workspace.model_path, device, max_new_tokens, MAX_LENGTH)
+        _LOG.info(f"inference engine: {engine.__class__.__name__}")
+
+        batch_size = batch_size or engine.get_default_batch_size()
+        _LOG.info(f"model loading time: {time.time() - t0:.2f}s")
+
+        if batch_size > sample_size:
+            batch_size = sample_size
+        _LOG.info(f"{batch_size=}")
+
+        # prepare seed data for clean consumption by formatron
+        seed_data = prepare_seed_for_formatron(seed_data, engine.tokenizer)
+        seeded_tgt_columns = seed_data.columns.to_list()
+        unseeded_tgt_columns = [c for c in tgt_text_columns if c not in seeded_tgt_columns]
+
+        total_tokenize_fn_time = 0
+        total_logits_processor_build_time = 0
+        total_generate_fn_time = 0
+
+        enforce_json_output = engine.supports_json_enforcing()
+        _LOG.info(f"{enforce_json_output=}")
+        formatron_vocab_processors = get_vocab_processors(is_peft_adapter)
+
+        if enforce_json_output and len(seeded_tgt_columns) == 0:
+            t0 = time.time()
+            formatter_builders = get_formatter_builders(size=batch_size, unseeded_fields=unseeded_tgt_columns)
+            engine.initialize_logits_processors(formatter_builders, formatron_vocab_processors)
+            total_logits_processor_build_time += time.time() - t0
+
+        # keep at most 500k samples in memory before decoding and writing to disk
+        buffer = FixedSizeSampleBuffer(capacity=500_000)
+
+        progress.update(completed=0, total=sample_size)
+        samples_processed = 0
+        while samples_processed < sample_size:
+            encoded_ctx_batch = encoded_ctx_data.iloc[samples_processed : samples_processed + batch_size]
+            sample_seed_batch = seed_data.iloc[samples_processed : samples_processed + batch_size]
+            ctx_batch = ctx_data.iloc[samples_processed : samples_processed + batch_size]
+            ctx_keys = ctx_batch[ctx_primary_key]
+
+            if enforce_json_output and len(seeded_tgt_columns) > 0:
+                t0 = time.time()
+                # some columns are seeded, so we need to create a new logits processor for each batch
+                formatter_builders = get_formatter_builders(
+                    seed_df=sample_seed_batch,
+                    unseeded_fields=unseeded_tgt_columns,
+                )
+                engine.initialize_logits_processors(formatter_builders, formatron_vocab_processors)
+                total_logits_processor_build_time += time.time() - t0
+
+            outputs, metrics = engine.generate(
+                encoded_ctx_batch["ctx"].tolist(),
+                sampling_temperature=sampling_temperature,
+                sampling_top_p=sampling_top_p,
+            )
+            total_tokenize_fn_time += metrics.tokenize_time
+            total_generate_fn_time += metrics.generate_time
+
+            buffer.add((outputs, ctx_keys, sample_seed_batch))
+            if buffer.is_full():
+                decoded_data = decode_buffered_samples(
+                    buffer, engine.tokenizer, tgt_text_columns, tgt_context_key, max_new_tokens
+                )
+                persist_data_part(
+                    decoded_data,
+                    output_path,
+                    f"{buffer.n_clears:06}.{0:06}",
+                )
+                buffer.clear()
+            progress.update(advance=len(ctx_batch))
+            samples_processed += len(ctx_batch)
+
+        if not buffer.is_empty():
+            decoded_data = decode_buffered_samples(
+                buffer, engine.tokenizer, tgt_text_columns, tgt_context_key, max_new_tokens
+            )
+            persist_data_part(
+                decoded_data,
+                output_path,
+                f"{buffer.n_clears:06}.{0:06}",
+            )
+            buffer.clear()
+        _LOG.info(f"{total_tokenize_fn_time=:.2f}s")
+        _LOG.info(f"{total_logits_processor_build_time=:.2f}s")
+        _LOG.info(f"{total_generate_fn_time=:.2f}s")
+        engine.cleanup()
+    _LOG.info(f"GENERATE_LANGUAGE finished in {time.time() - t0_:.2f}s")

mostlyai/engine/_language/lstm.py

@@ -0,0 +1,139 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import torch
+from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModelForCausalLM, AutoModel, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutput
+import torch.nn as nn
+
+_LOG = logging.getLogger(__name__)
+
+
+class LSTMFromScratchConfig(PretrainedConfig):
+    model_type = model_id = "MOSTLY_AI/LSTMFromScratch-3m"
+
+    def __init__(
+        self,
+        vocab_size: int | None = None,
+        embedding_size: int = 256,
+        hidden_size: int = 256,
+        num_layers: int = 1,
+        dropout: float = 0.25,
+        with_dp: bool = False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.embedding_size = embedding_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.with_dp = with_dp
+        super().__init__(**kwargs)
+
+
+class LSTMFromScratchLMHeadModel(PreTrainedModel, GenerationMixin):
+    config_class = LSTMFromScratchConfig
+
+    def __init__(self, config: LSTMFromScratchConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embedding = nn.Embedding(self.config.vocab_size, self.config.embedding_size)
+        self.dropout = nn.Dropout(self.config.dropout)
+        if self.config.with_dp:
+            from opacus.layers import DPLSTM
+
+            lstm_cls = DPLSTM
+        else:
+            lstm_cls = nn.LSTM
+        self.lstm = lstm_cls(
+            input_size=self.config.embedding_size,
+            hidden_size=self.config.hidden_size,
+            num_layers=self.config.num_layers,
+            dropout=self.config.dropout if self.config.num_layers > 1 else 0.0,
+            batch_first=True,
+        )
+        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size)
+        self.loss_fn = nn.CrossEntropyLoss()
+
+        # this will be filled by left_to_right_padding() during the generation
+        self.pad_token_id = None
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        **kwargs,
+    ) -> CausalLMOutput:
+        lengths = attention_mask.sum(dim=1)
+        embeddings = self.embedding(input_ids)
+        embeddings = self.dropout(embeddings)
+
+        # (DP)LSTM layers without pack_padded_sequence/pad_packed_sequence
+        lstm_outputs, _ = self.lstm(embeddings)
+
+        logits = self.lm_head(lstm_outputs)
+
+        loss = None
+        if labels is not None:
+            labels = labels[:, 1:].contiguous()
+            shifted_prediction_scores = logits[:, :-1, :].contiguous()
+            loss = self.loss_fn(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+        else:
+            # overwrite the logit of the last time step with the logit of the actual last token
+            # so that Hugging Face Transformers' generate() will sample on the right probabilities
+            logits[:, -1, :] = torch.stack([logits[i, length - 1, :] for i, length in enumerate(lengths)])
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
+    ) -> dict[str, torch.Tensor]:
+        """
+        This function is mandatory so that the model is able to use the Hugging Face `.generate()` method.
+        Since `.generate()` works with left-padded sequences but the model is trained with right-padded sequences,
+        we need to convert the padding side here to make it work properly.
+        """
+        lengths = attention_mask.sum(dim=1)
+        return {
+            "input_ids": self.left_to_right_padding(input_ids, lengths),
+            "attention_mask": attention_mask,
+        }
+
+    def left_to_right_padding(self, left_padded_tensors: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
+        batch_size, max_length = left_padded_tensors.size()
+        indices = torch.nonzero(lengths < max_length)
+        if len(indices) == 0:
+            # none of the samples are padded, so we can just return them as they are
+            return left_padded_tensors
+        else:
+            if self.pad_token_id is None:
+                # get the pad token id from the first padded sample
+                self.pad_token_id = left_padded_tensors[indices[0], -1].item()
+            right_padded_tensors = torch.full_like(left_padded_tensors, self.pad_token_id)
+            for i in range(batch_size):
+                right_padded_tensors[i, : lengths[i]] = left_padded_tensors[i, max_length - lengths[i] :]
+            return right_padded_tensors
+
+
+def register_mostly_lstm_model():
+    # register the model so that we can load it with `AutoModelForCausalLM.from_pretrained()` later
+    AutoConfig.register(LSTMFromScratchConfig.model_id, LSTMFromScratchConfig)
+    AutoModel.register(LSTMFromScratchConfig, LSTMFromScratchLMHeadModel)
+    AutoModelForCausalLM.register(LSTMFromScratchConfig, LSTMFromScratchLMHeadModel)

mostlyai/engine/_language/tokenizer_utils.py

@@ -0,0 +1,144 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any
+from collections.abc import Mapping, Iterator
+
+from transformers import DataCollatorForLanguageModeling, BatchEncoding, PreTrainedTokenizerFast, LlamaTokenizerFast
+from transformers.data.data_collator import pad_without_fast_tokenizer_warning, _torch_collate_batch
+
+
+#################
+### TOKENIZER ###
+#################
+
+
+def train_tokenizer(training_iterator: Iterator | list | None = None, tokenizer_kwargs=None):
+    if tokenizer_kwargs is None:
+        tokenizer_kwargs = {}
+    from tokenizers import Tokenizer, decoders
+    from tokenizers.models import BPE
+    from tokenizers.normalizers import Replace
+    from tokenizers.pre_tokenizers import Metaspace, Punctuation, Sequence, Split
+    from tokenizers.trainers import BpeTrainer
+
+    special_tokens = {
+        "unk_token": "<unk>",
+        "pad_token": "<pad>",
+        "bos_token": "<s>",
+        "eos_token": "</s>",
+    }
+    SPECIAL_TOKENS = list(special_tokens.values())
+    NEW_LINE_VALUE = "\n"
+    NEW_LINE_SYMBOL = "\u240a"  # https://www.fileformat.info/info/unicode/char/240a/index.htm
+    MIN_FREQ_MERGE = 20
+    VOCAB_SIZE = 5000
+
+    # Builds a BPE raw_tokenizer, and optionally trains it based on provided text
+    training_iterator = training_iterator or []  # allow easy training skip
+    raw_tokenizer = Tokenizer(BPE(unk_token=special_tokens["unk_token"]))
+    trainer = BpeTrainer(
+        special_tokens=SPECIAL_TOKENS,
+        min_frequency=MIN_FREQ_MERGE,
+        vocab_size=VOCAB_SIZE,
+        show_progress=False,
+    )
+    raw_tokenizer.normalizer = Replace(NEW_LINE_VALUE, NEW_LINE_SYMBOL)
+    raw_tokenizer.pre_tokenizer = Sequence(
+        [
+            Metaspace(),
+            Split(pattern=NEW_LINE_SYMBOL, behavior="isolated"),
+            Punctuation(),
+        ]
+    )
+    raw_tokenizer.decoder = decoders.Sequence(
+        [
+            decoders.Metaspace(),
+            decoders.Replace(NEW_LINE_SYMBOL, NEW_LINE_VALUE),
+        ]
+    )
+    raw_tokenizer.train_from_iterator(iterator=training_iterator, trainer=trainer)
+    tokenizer = LlamaTokenizerFast(tokenizer_object=raw_tokenizer, **special_tokens, **tokenizer_kwargs)
+    return tokenizer
+
+
+def tokenize_fn(
+    text: dict[str, str] | dict[str, list[str]] | list[str],
+    tokenizer: PreTrainedTokenizerFast,
+    text_key: str | None = None,
+    return_tensors: str | None = None,
+    padding: bool | str = True,
+    truncation: bool = True,
+    add_bos_token: bool = True,
+    add_eos_token: bool = True,
+    max_length: int = 1024,
+) -> BatchEncoding:
+    if text_key:
+        text = text[text_key]
+    # make sure the tokenizer is configured as expected
+    if getattr(tokenizer, "add_bos_token", False) or getattr(tokenizer, "add_eos_token", False):
+        raise RuntimeError("Tokenizer must be configured as add_bos_token=False and add_eos_token=False")
+    if tokenizer.bos_token is None or tokenizer.eos_token is None:
+        raise RuntimeError("Tokenizer must have bos_token and eos_token set")
+    prefix = tokenizer.bos_token if add_bos_token else ""
+    suffix = tokenizer.eos_token if add_eos_token else ""
+    # NOTE: here we add bos/eos tokens before truncation and padding,
+    # which means that they may be truncated for long sequences
+    if isinstance(text, str):
+        text = f"{prefix}{text}{suffix}"
+    else:
+        for i, t in enumerate(text):
+            text[i] = f"{prefix}{t}{suffix}"
+    tokenized_content = tokenizer(
+        text,
+        padding=padding,
+        truncation=truncation,
+        max_length=max_length,
+        return_tensors=return_tensors,
+    )
+    return tokenized_content
+
+
+#####################
+### DATA COLLATOR ###
+#####################
+
+
+@dataclass
+class MostlyDataCollatorForLanguageModeling(DataCollatorForLanguageModeling):
+    def torch_call(self, examples: list[list[int] | Any | dict[str, Any]]) -> dict[str, Any]:
+        """
+        A variation of the original `DataCollatorForLanguageModeling.torch_call` method.
+
+        This method can mask tokens based on the attention mask, so that bos and eos tokens will not be masked
+        even if they are identical to pad token.
+        If attention mask is not provided, it will fall back to masking pad tokens.
+        """
+        if isinstance(examples[0], Mapping):
+            batch = pad_without_fast_tokenizer_warning(
+                self.tokenizer, examples, return_tensors="pt", pad_to_multiple_of=None
+            )
+        else:
+            batch = {"input_ids": _torch_collate_batch(examples, self.tokenizer, pad_to_multiple_of=None)}
+
+        labels = batch["input_ids"].clone()
+        attention_mask = batch.get("attention_mask", None)
+        if attention_mask is not None:
+            labels[(attention_mask == 0)] = -100
+        else:
+            if self.tokenizer.pad_token_id is not None:
+                labels[labels == self.tokenizer.pad_token_id] = -100
+        batch["labels"] = labels
+        return batch