Remove byte tokenizer and add config args to switch between byte/patch packing (#68)
Some checks are pending
Lint with Black / lint (push) Waiting to run
Lint with isort / lint (push) Waiting to run

Summary:

Test Plan:

```
python -m bytelatent.train config=../internal-blt/configs/entropy_model.yaml logging.wandb=null checkpoint.dump.every=1000 checkpoint.eval.every=100000 eval=null

pytest bytelatent/
```
This commit is contained in:
Pedro Rodriguez 2025-02-25 11:10:59 -08:00 committed by GitHub
parent ff36aa8642
commit aeb95f12a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 23 additions and 45 deletions

View file

@ -1,4 +1,5 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
from enum import Enum
from typing import Any
import numpy as np
@ -12,6 +13,11 @@ from bytelatent.data.iterators.abstract_iterator import (
from bytelatent.data.iterators.sampling_iterator import SamplingIteratorState
class PackingMode(str, Enum):
BYTES = "bytes"
PATCHING = "patching"
class PackingArgs(BaseModel):
model_config = ConfigDict(extra="forbid")
batch_size: int
@ -20,7 +26,7 @@ class PackingArgs(BaseModel):
max_length: int | None
pad_to_max_length: bool
enable_byte_ngrams: bool
tokenizer_name: str
packing_mode: PackingMode
class PackingIteratorState(PydanticIteratorState):
@ -155,10 +161,12 @@ class PackingIterator(StatefulIterator[Batch, PackingIteratorState]):
)
def create_iter(self):
if self.packing_args.tokenizer_name == "bytes":
if self.packing_args.packing_mode == PackingMode.BYTES:
return self._create_iter_from_bytes()
else:
elif self.packing_args.packing_mode == PackingMode.PATCHING:
return self._create_iter_from_patch_lengths()
else:
raise ValueError(f"Invalid patching mode: {self.packing_args.packing_mode}")
def _create_iter_from_bytes(self):
sequence_iter = self.sequence_iterator.create_iter()