tacl-llm 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (48) hide show
  1. tacl_llm-2.0.0/PKG-INFO +26 -0
  2. tacl_llm-2.0.0/README.md +17 -0
  3. tacl_llm-2.0.0/pyproject.toml +46 -0
  4. tacl_llm-2.0.0/setup.cfg +4 -0
  5. tacl_llm-2.0.0/src/tacl/__init__.py +29 -0
  6. tacl_llm-2.0.0/src/tacl/auto.py +51 -0
  7. tacl_llm-2.0.0/src/tacl/configuration_tacl.py +132 -0
  8. tacl_llm-2.0.0/src/tacl/data/__init__.py +14 -0
  9. tacl_llm-2.0.0/src/tacl/data/base.py +33 -0
  10. tacl_llm-2.0.0/src/tacl/data/pretrain_dataset.py +171 -0
  11. tacl_llm-2.0.0/src/tacl/data/sft_dataset.py +200 -0
  12. tacl_llm-2.0.0/src/tacl/data/streaming.py +52 -0
  13. tacl_llm-2.0.0/src/tacl/exceptions.py +9 -0
  14. tacl_llm-2.0.0/src/tacl/generation/__init__.py +7 -0
  15. tacl_llm-2.0.0/src/tacl/generation/generate.py +223 -0
  16. tacl_llm-2.0.0/src/tacl/generation/streamer.py +131 -0
  17. tacl_llm-2.0.0/src/tacl/modeling_tacl.py +389 -0
  18. tacl_llm-2.0.0/src/tacl/modules/__init__.py +17 -0
  19. tacl_llm-2.0.0/src/tacl/modules/layer.py +38 -0
  20. tacl_llm-2.0.0/src/tacl/modules/linear_attention.py +72 -0
  21. tacl_llm-2.0.0/src/tacl/modules/quantization.py +61 -0
  22. tacl_llm-2.0.0/src/tacl/modules/sparsemax.py +27 -0
  23. tacl_llm-2.0.0/src/tacl/modules/swiglu_ffn.py +23 -0
  24. tacl_llm-2.0.0/src/tacl/modules/ternary_linear.py +103 -0
  25. tacl_llm-2.0.0/src/tacl/modules/thought_block.py +65 -0
  26. tacl_llm-2.0.0/src/tacl/presets.py +80 -0
  27. tacl_llm-2.0.0/src/tacl/scripts/__init__.py +0 -0
  28. tacl_llm-2.0.0/src/tacl/scripts/build_vocab.py +108 -0
  29. tacl_llm-2.0.0/src/tacl/scripts/generate.py +67 -0
  30. tacl_llm-2.0.0/src/tacl/scripts/import_vocab.py +46 -0
  31. tacl_llm-2.0.0/src/tacl/scripts/smoke_check.py +303 -0
  32. tacl_llm-2.0.0/src/tacl/scripts/train.py +106 -0
  33. tacl_llm-2.0.0/src/tacl/serialization.py +227 -0
  34. tacl_llm-2.0.0/src/tacl/tokenization_tacl.py +223 -0
  35. tacl_llm-2.0.0/src/tacl/training/__init__.py +11 -0
  36. tacl_llm-2.0.0/src/tacl/training/callbacks.py +59 -0
  37. tacl_llm-2.0.0/src/tacl/training/chl.py +183 -0
  38. tacl_llm-2.0.0/src/tacl/training/logging_utils.py +99 -0
  39. tacl_llm-2.0.0/src/tacl/training/trainer.py +306 -0
  40. tacl_llm-2.0.0/src/tacl/utils/__init__.py +4 -0
  41. tacl_llm-2.0.0/src/tacl/utils/device.py +16 -0
  42. tacl_llm-2.0.0/src/tacl/utils/seed.py +11 -0
  43. tacl_llm-2.0.0/src/tacl_llm.egg-info/PKG-INFO +26 -0
  44. tacl_llm-2.0.0/src/tacl_llm.egg-info/SOURCES.txt +46 -0
  45. tacl_llm-2.0.0/src/tacl_llm.egg-info/dependency_links.txt +1 -0
  46. tacl_llm-2.0.0/src/tacl_llm.egg-info/entry_points.txt +4 -0
  47. tacl_llm-2.0.0/src/tacl_llm.egg-info/requires.txt +24 -0
  48. tacl_llm-2.0.0/src/tacl_llm.egg-info/top_level.txt +1 -0
@@ -0,0 +1,26 @@
1
+ Metadata-Version: 2.4
2
+ Name: tacl-llm
3
+ Version: 2.0.0
4
+ Summary: Ternary Additive Contrastive LLM — CPU-friendly integer neural network
5
+ Author: TACL Contributors
6
+ License-Expression: MIT
7
+ Requires-Python: >=3.10
8
+ Requires-Dist: numpy>=1.24
9
+ Requires-Dist: torch>=2.2
10
+ Requires-Dist: safetensors>=0.4
11
+ Requires-Dist: tokenizers>=0.19
12
+ Requires-Dist: huggingface_hub>=0.23
13
+ Requires-Dist: tqdm
14
+ Requires-Dist: pyyaml
15
+ Provides-Extra: colab
16
+ Requires-Dist: ipywidgets; extra == "colab"
17
+ Requires-Dist: matplotlib; extra == "colab"
18
+ Provides-Extra: logging
19
+ Requires-Dist: tensorboard; extra == "logging"
20
+ Provides-Extra: gpu
21
+ Provides-Extra: datasets
22
+ Requires-Dist: datasets>=2.14; extra == "datasets"
23
+ Provides-Extra: dev
24
+ Requires-Dist: pytest>=7; extra == "dev"
25
+ Requires-Dist: pytest-cov; extra == "dev"
26
+ Requires-Dist: datasets>=2.14; extra == "dev"
@@ -0,0 +1,17 @@
1
+ # TACL — Ternary Additive Contrastive LLM
2
+
3
+ CPU-friendly integer neural network. Not a Transformer — uses CHL (Equilibrium Propagation), sparsemax, linear attention, and ternary weights.
4
+
5
+ ## Quick Install
6
+
7
+ ```bash
8
+ pip install tacl-llm
9
+ ```
10
+
11
+ ## Documentation
12
+
13
+ See [ts/00-overview-and-index.md](ts/00-overview-and-index.md) for the full specification, and [wiki/](wiki/) for research wiki and verification reports.
14
+
15
+ ## License
16
+
17
+ MIT
@@ -0,0 +1,46 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "tacl-llm"
7
+ version = "2.0.0"
8
+ description = "Ternary Additive Contrastive LLM — CPU-friendly integer neural network"
9
+ requires-python = ">=3.10"
10
+ authors = [{ name = "TACL Contributors" }]
11
+ license = "MIT"
12
+
13
+ dependencies = [
14
+ "numpy>=1.24",
15
+ "torch>=2.2",
16
+ "safetensors>=0.4",
17
+ "tokenizers>=0.19",
18
+ "huggingface_hub>=0.23",
19
+ "tqdm",
20
+ "pyyaml",
21
+ ]
22
+
23
+ [project.optional-dependencies]
24
+ colab = ["ipywidgets", "matplotlib"]
25
+ logging = ["tensorboard"]
26
+ gpu = []
27
+ datasets = ["datasets>=2.14"]
28
+ dev = ["pytest>=7", "pytest-cov", "datasets>=2.14"]
29
+
30
+ [project.scripts]
31
+ tacl-train = "tacl.scripts.train:main"
32
+ tacl-generate = "tacl.scripts.generate:main"
33
+ tacl-smoke-check = "tacl.scripts.smoke_check:run_smoke_check"
34
+
35
+ [tool.setuptools.packages.find]
36
+ where = ["src"]
37
+
38
+ [tool.setuptools.package-data]
39
+ tacl = ["py.typed"]
40
+
41
+ [tool.pytest.ini_options]
42
+ testpaths = ["tests"]
43
+ pythonpath = ["src"]
44
+ filterwarnings = [
45
+ "ignore::tacl.exceptions.TACLCaveatWarning",
46
+ ]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,29 @@
1
+ from .configuration_tacl import TACLConfig
2
+ from .modeling_tacl import TACLModel, TACLForCausalLM
3
+ from .tokenization_tacl import TACLTokenizer
4
+ from .training.trainer import TACLTrainer, TACLTrainingArguments
5
+ from .auto import TACLAutoModel, TACLAutoForCausalLM, TACLAutoConfig, TACLAutoTokenizer
6
+ from .presets import PRESETS, list_presets
7
+ from .generation.streamer import TACLTextStreamer, TACLTextIteratorStreamer
8
+ from .data import PretrainPackedDataset, SFTDataset
9
+
10
+ __all__ = [
11
+ "TACLConfig",
12
+ "TACLModel",
13
+ "TACLForCausalLM",
14
+ "TACLTokenizer",
15
+ "TACLTrainer",
16
+ "TACLTrainingArguments",
17
+ "TACLAutoModel",
18
+ "TACLAutoForCausalLM",
19
+ "TACLAutoConfig",
20
+ "TACLAutoTokenizer",
21
+ "PRESETS",
22
+ "list_presets",
23
+ "TACLTextStreamer",
24
+ "TACLTextIteratorStreamer",
25
+ "PretrainPackedDataset",
26
+ "SFTDataset",
27
+ ]
28
+
29
+ __version__ = "2.0.0"
@@ -0,0 +1,51 @@
1
+ from __future__ import annotations
2
+
3
+ from pathlib import Path
4
+ from typing import Any, Optional
5
+
6
+ from .configuration_tacl import TACLConfig
7
+ from .modeling_tacl import TACLForCausalLM, TACLModel
8
+ from .tokenization_tacl import TACLTokenizer
9
+ from .presets import PRESETS
10
+
11
+
12
+ class TACLAutoModel:
13
+ @classmethod
14
+ def from_pretrained(cls, path_or_repo: str) -> TACLModel:
15
+ config = TACLConfig.from_pretrained(path_or_repo)
16
+ model = TACLModel(config)
17
+ from .serialization import load_state_dict
18
+ load_state_dict(model, path_or_repo)
19
+ return model
20
+
21
+ @classmethod
22
+ def from_preset(cls, name: str, **overrides: Any) -> TACLModel:
23
+ if name not in PRESETS:
24
+ raise KeyError(f"Unknown preset {name!r}. Available: {list(PRESETS.keys())}")
25
+ config = PRESETS[name](**overrides)
26
+ return TACLModel(config)
27
+
28
+
29
+ class TACLAutoForCausalLM:
30
+ @classmethod
31
+ def from_pretrained(cls, path_or_repo: str) -> TACLForCausalLM:
32
+ return TACLForCausalLM.from_pretrained(path_or_repo)
33
+
34
+ @classmethod
35
+ def from_preset(cls, name: str, **overrides: Any) -> TACLForCausalLM:
36
+ return TACLForCausalLM.from_preset(name, **overrides)
37
+
38
+
39
+ class TACLAutoConfig:
40
+ @classmethod
41
+ def from_pretrained(cls, path_or_repo: str) -> TACLConfig:
42
+ return TACLConfig.from_pretrained(path_or_repo)
43
+
44
+
45
+ class TACLAutoTokenizer:
46
+ @classmethod
47
+ def from_pretrained(cls, path_or_repo: str, use_hf: bool = False, **kwargs: Any) -> TACLTokenizer:
48
+ from pathlib import Path as _Path
49
+ if use_hf or (not _Path(path_or_repo).exists() and "/" in path_or_repo):
50
+ return TACLTokenizer.from_pretrained_hf(path_or_repo, **kwargs)
51
+ return TACLTokenizer.from_pretrained(path_or_repo)
@@ -0,0 +1,132 @@
1
+ from __future__ import annotations
2
+
3
+ import json
4
+ import warnings
5
+ from dataclasses import dataclass, asdict
6
+ from pathlib import Path
7
+ from typing import Any
8
+
9
+ from .exceptions import TACLCaveatWarning, warn_caveat
10
+
11
+
12
+ _CAVEATS = {
13
+ "chl_validated_on": "MLP proxy only, full-stack integration PENDING (verification-report.md §2.1, §2.3)",
14
+ "language_modeling_quality": "unmeasured — no perplexity/accuracy benchmark yet",
15
+ "context_decay": "S_t has no decay mechanism, uniform weighting over all past tokens (see final-architecture.md §2.2 'Known limitation')",
16
+ "max_safe_context_int32": 4096,
17
+ }
18
+
19
+ _TACL_VERSION = "2.0.0"
20
+
21
+
22
+ @dataclass
23
+ class TACLConfig:
24
+ tacl_version: str = "2.0"
25
+ vocab_size: int = 8192
26
+ hidden_size: int = 512
27
+ ffn_size: int = 2048
28
+ num_layers: int = 8
29
+ max_position_embeddings: int = 4096
30
+
31
+ ternary_sparsity_target: float = 0.5
32
+ quant_format: str = "bfp_int8"
33
+ fixed_point_qf: int = 12
34
+ stochastic_rounding: bool = True
35
+
36
+ attn_kernel: str = "elu1"
37
+ use_thought_loop: bool = False
38
+ thought_max_iters: int = 8
39
+ thought_epsilon_qf_units: int = 1
40
+
41
+ learning_rule: str = "chl"
42
+ chl_alpha: float = 0.05
43
+
44
+ tie_word_embeddings: bool = True
45
+ pad_token_id: int = 0
46
+ bos_token_id: int = 1
47
+ eos_token_id: int = 2
48
+ unk_token_id: int = 3
49
+
50
+ use_fixed_residual: bool = False
51
+
52
+ _caveats: dict[str, str] | None = None
53
+
54
+ def __post_init__(self) -> None:
55
+ if self.learning_rule not in ("chl",):
56
+ raise ValueError(
57
+ f"learning_rule={self.learning_rule!r} не поддерживается: DFA отвергнут "
58
+ "численно (cos~0.19 c BP, +332% loss при обучении, см. verification-report.md §1.3/§2.2)."
59
+ )
60
+ if self.hidden_size % 2 != 0:
61
+ raise ValueError("hidden_size должен быть чётным (используется в d/2 расчётах FLOPs)")
62
+ if self.max_position_embeddings > 4096:
63
+ warnings.warn(
64
+ f"max_position_embeddings={self.max_position_embeddings} > 4096: "
65
+ "int32 overflow risk in linear attention state. "
66
+ "See verification-report.md §3.2 for details.",
67
+ stacklevel=2,
68
+ )
69
+ if self._caveats is None:
70
+ self._caveats = dict(_CAVEATS)
71
+
72
+ def to_dict(self) -> dict[str, Any]:
73
+ d = {}
74
+ for k, v in asdict(self).items():
75
+ if k == "_caveats" and not v:
76
+ continue
77
+ d[k] = v
78
+ d["auto_map"] = {
79
+ "AutoConfig": "tacl.configuration_tacl.TACLConfig",
80
+ "AutoModel": "tacl.auto.TACLAutoModel",
81
+ "AutoModelForCausalLM": "tacl.auto.TACLAutoForCausalLM",
82
+ }
83
+ return d
84
+
85
+ def to_json_string(self) -> str:
86
+ return json.dumps(self.to_dict(), indent=2, ensure_ascii=False)
87
+
88
+ @classmethod
89
+ def from_dict(cls, d: dict[str, Any]) -> "TACLConfig":
90
+ return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
91
+
92
+ @classmethod
93
+ def from_json_file(cls, path: str | Path) -> "TACLConfig":
94
+ with open(path, "r", encoding="utf-8") as f:
95
+ return cls.from_dict(json.load(f))
96
+
97
+ @classmethod
98
+ def _version_tuple(cls, v: str) -> tuple[int, ...]:
99
+ parts = v.split(".")
100
+ return tuple(int(p) if p.isdigit() else 0 for p in parts)
101
+
102
+ def save_pretrained(self, directory: str | Path) -> None:
103
+ directory = Path(directory)
104
+ directory.mkdir(parents=True, exist_ok=True)
105
+ self.tacl_version = _TACL_VERSION
106
+ with open(directory / "config.json", "w", encoding="utf-8") as f:
107
+ f.write(self.to_json_string())
108
+
109
+ @classmethod
110
+ def from_pretrained(cls, directory_or_repo: str | Path) -> "TACLConfig":
111
+ path = Path(directory_or_repo)
112
+ if not path.exists():
113
+ from huggingface_hub import hf_hub_download
114
+ config_path = hf_hub_download(str(directory_or_repo), filename="config.json")
115
+ else:
116
+ config_path = path / "config.json"
117
+ config = cls.from_json_file(config_path)
118
+
119
+ saved_version = config.tacl_version
120
+ if saved_version and cls._version_tuple(saved_version) != cls._version_tuple(_TACL_VERSION):
121
+ warnings.warn(
122
+ f"TACL version mismatch: saved model uses tacl_version={saved_version}, "
123
+ f"current library is {_TACL_VERSION}. "
124
+ "Upgrade/downgrade may be needed.",
125
+ stacklevel=2,
126
+ )
127
+
128
+ if config._caveats:
129
+ for key, val in config._caveats.items():
130
+ warn_caveat(f"{key}: {val}")
131
+
132
+ return config
@@ -0,0 +1,14 @@
1
+ from .base import TACLDataset
2
+ from .pretrain_dataset import PretrainPackedDataset, pretrain_collate_fn
3
+ from .sft_dataset import SFTDataset, DEFAULT_TACL_CHAT_TEMPLATE, sft_collate_fn
4
+ from .streaming import StreamingTextDataset
5
+
6
+ __all__ = [
7
+ "TACLDataset",
8
+ "PretrainPackedDataset",
9
+ "pretrain_collate_fn",
10
+ "SFTDataset",
11
+ "DEFAULT_TACL_CHAT_TEMPLATE",
12
+ "sft_collate_fn",
13
+ "StreamingTextDataset",
14
+ ]
@@ -0,0 +1,33 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Any, Optional, Union
4
+
5
+ import torch
6
+ from torch.utils.data import Dataset, IterableDataset
7
+
8
+
9
+ class TACLDataset(Dataset):
10
+ def __init__(
11
+ self,
12
+ source: Union[str, list[dict]],
13
+ tokenizer: Any,
14
+ max_length: int = 512,
15
+ ):
16
+ self._tokenizer = tokenizer
17
+ self._max_length = max_length
18
+ self._samples: list[dict[str, torch.Tensor]] = []
19
+ self._load(source)
20
+
21
+ def _load(self, source: Union[str, list[dict]]) -> None:
22
+ raise NotImplementedError
23
+
24
+ def __len__(self) -> int:
25
+ return len(self._samples)
26
+
27
+ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
28
+ return self._samples[idx]
29
+
30
+ def _get_eos_id(self) -> int:
31
+ if self._tokenizer._config is not None and hasattr(self._tokenizer._config, 'eos_token_id'):
32
+ return self._tokenizer._config.eos_token_id
33
+ return 2
@@ -0,0 +1,171 @@
1
+ from __future__ import annotations
2
+
3
+ import csv
4
+ import glob
5
+ import logging
6
+ import time
7
+ from pathlib import Path
8
+ from typing import Any, Iterable, Iterator, Optional, Union
9
+
10
+ import torch
11
+ from torch.utils.data import DataLoader, IterableDataset
12
+
13
+ logger = logging.getLogger(__name__)
14
+
15
+
16
+ def _get_eos_id(tokenizer: Any) -> int:
17
+ if hasattr(tokenizer, '_config') and tokenizer._config is not None and hasattr(tokenizer._config, 'eos_token_id'):
18
+ return tokenizer._config.eos_token_id
19
+ return 2
20
+
21
+
22
+ class PretrainPackedDataset(IterableDataset):
23
+ """
24
+ Tokenises a text stream and packs into contiguous blocks of seq_len,
25
+ without padding inside each block (standard GPT-style "packing").
26
+
27
+ Supports:
28
+ - Local text files (.txt, .csv, .parquet, glob patterns)
29
+ - HF Datasets (requires datasets package)
30
+ - Arbitrary Python iterator of strings
31
+
32
+ Shuffle is done via a fixed-size reservoir buffer, not full shuffle.
33
+ """
34
+
35
+ def __init__(
36
+ self,
37
+ source: Union[str, Iterable[str]],
38
+ tokenizer: Any,
39
+ seq_len: int = 1024,
40
+ text_column: str = "text",
41
+ hf_dataset_config: Optional[str] = None,
42
+ hf_split: str = "train",
43
+ shuffle_buffer_size: int = 10000,
44
+ add_eos_between_docs: bool = True,
45
+ ):
46
+ self._source = source
47
+ self._tokenizer = tokenizer
48
+ self._seq_len = seq_len
49
+ self._text_column = text_column
50
+ self._hf_dataset_config = hf_dataset_config
51
+ self._hf_split = hf_split
52
+ self._shuffle_buffer_size = shuffle_buffer_size
53
+ self._add_eos = add_eos_between_docs
54
+
55
+ self._doc_iter: Optional[Iterator[str]] = None
56
+ self._token_buffer: list[int] = []
57
+ self._docs_processed = 0
58
+ self._tokens_processed = 0
59
+ self._start_time = time.time()
60
+
61
+ def _get_iterator(self) -> Iterator[str]:
62
+ if isinstance(self._source, str):
63
+ matched_files = glob.glob(self._source)
64
+ if matched_files:
65
+ for path_str in matched_files:
66
+ path = Path(path_str)
67
+ ext = path.suffix.lower()
68
+ if ext == ".csv":
69
+ with open(path, "r", encoding="utf-8", newline="") as f:
70
+ reader = csv.DictReader(f)
71
+ for row in reader:
72
+ yield row.get(self._text_column, "")
73
+ elif ext == ".parquet":
74
+ try:
75
+ import pandas as pd
76
+ df = pd.read_parquet(path)
77
+ for val in df[self._text_column].tolist():
78
+ yield str(val)
79
+ except ImportError:
80
+ raise ImportError("Install pandas and pyarrow for .parquet support")
81
+ else:
82
+ with open(path, "r", encoding="utf-8", errors="ignore") as f:
83
+ yield from f
84
+ else:
85
+ try:
86
+ from datasets import load_dataset
87
+ ds = load_dataset(
88
+ self._source,
89
+ self._hf_dataset_config,
90
+ split=self._hf_split,
91
+ streaming=True,
92
+ )
93
+ for example in ds:
94
+ yield example[self._text_column]
95
+ except ImportError:
96
+ raise ImportError(
97
+ f"To use HF datasets, install with: pip install tacl[datasets]"
98
+ )
99
+ except Exception:
100
+ with open(self._source, "r", encoding="utf-8", errors="ignore") as f:
101
+ yield from f
102
+ elif isinstance(self._source, Iterable):
103
+ yield from self._source
104
+ else:
105
+ raise TypeError(f"Unsupported source type: {type(self._source)}")
106
+
107
+ def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
108
+ self._token_buffer = []
109
+ self._docs_processed = 0
110
+ self._tokens_processed = 0
111
+ self._start_time = time.time()
112
+ doc_iter = self._get_iterator()
113
+
114
+ buf: list[dict[str, torch.Tensor]] = []
115
+
116
+ for doc_text in doc_iter:
117
+ doc_text = doc_text.strip()
118
+ if not doc_text:
119
+ continue
120
+
121
+ ids = self._tokenizer.encode(doc_text)
122
+ self._docs_processed += 1
123
+
124
+ if self._add_eos:
125
+ eos_id = _get_eos_id(self._tokenizer)
126
+ ids = ids + [eos_id]
127
+
128
+ self._token_buffer.extend(ids)
129
+
130
+ while len(self._token_buffer) >= self._seq_len:
131
+ chunk = self._token_buffer[:self._seq_len]
132
+ self._token_buffer = self._token_buffer[self._seq_len:]
133
+ item = {
134
+ "input_ids": torch.tensor(chunk, dtype=torch.long),
135
+ "labels": torch.tensor(chunk, dtype=torch.long),
136
+ }
137
+ self._tokens_processed += self._seq_len
138
+
139
+ if self._shuffle_buffer_size > 0:
140
+ buf.append(item)
141
+ if len(buf) >= self._shuffle_buffer_size:
142
+ idx = torch.randint(0, len(buf), (1,)).item()
143
+ yield buf.pop(idx)
144
+ else:
145
+ yield item
146
+
147
+ if self._docs_processed % 1000 == 0:
148
+ elapsed = time.time() - self._start_time
149
+ tok_per_sec = self._tokens_processed / max(elapsed, 1e-6)
150
+ logger.info(
151
+ f"PretrainPackedDataset: {self._docs_processed} docs, "
152
+ f"{self._tokens_processed} tokens, {tok_per_sec:.0f} tok/s"
153
+ )
154
+
155
+ while buf:
156
+ idx = torch.randint(0, len(buf), (1,)).item()
157
+ yield buf.pop(idx)
158
+
159
+ remainder = self._token_buffer[:self._seq_len]
160
+ if len(remainder) > 0:
161
+ remainder = remainder + [0] * (self._seq_len - len(remainder))
162
+ yield {
163
+ "input_ids": torch.tensor(remainder, dtype=torch.long),
164
+ "labels": torch.tensor(remainder, dtype=torch.long),
165
+ }
166
+
167
+
168
+ def pretrain_collate_fn(batch: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
169
+ input_ids = torch.stack([item["input_ids"] for item in batch])
170
+ labels = torch.stack([item["labels"] for item in batch])
171
+ return {"input_ids": input_ids, "labels": labels}