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.
- tacl_llm-2.0.0/PKG-INFO +26 -0
- tacl_llm-2.0.0/README.md +17 -0
- tacl_llm-2.0.0/pyproject.toml +46 -0
- tacl_llm-2.0.0/setup.cfg +4 -0
- tacl_llm-2.0.0/src/tacl/__init__.py +29 -0
- tacl_llm-2.0.0/src/tacl/auto.py +51 -0
- tacl_llm-2.0.0/src/tacl/configuration_tacl.py +132 -0
- tacl_llm-2.0.0/src/tacl/data/__init__.py +14 -0
- tacl_llm-2.0.0/src/tacl/data/base.py +33 -0
- tacl_llm-2.0.0/src/tacl/data/pretrain_dataset.py +171 -0
- tacl_llm-2.0.0/src/tacl/data/sft_dataset.py +200 -0
- tacl_llm-2.0.0/src/tacl/data/streaming.py +52 -0
- tacl_llm-2.0.0/src/tacl/exceptions.py +9 -0
- tacl_llm-2.0.0/src/tacl/generation/__init__.py +7 -0
- tacl_llm-2.0.0/src/tacl/generation/generate.py +223 -0
- tacl_llm-2.0.0/src/tacl/generation/streamer.py +131 -0
- tacl_llm-2.0.0/src/tacl/modeling_tacl.py +389 -0
- tacl_llm-2.0.0/src/tacl/modules/__init__.py +17 -0
- tacl_llm-2.0.0/src/tacl/modules/layer.py +38 -0
- tacl_llm-2.0.0/src/tacl/modules/linear_attention.py +72 -0
- tacl_llm-2.0.0/src/tacl/modules/quantization.py +61 -0
- tacl_llm-2.0.0/src/tacl/modules/sparsemax.py +27 -0
- tacl_llm-2.0.0/src/tacl/modules/swiglu_ffn.py +23 -0
- tacl_llm-2.0.0/src/tacl/modules/ternary_linear.py +103 -0
- tacl_llm-2.0.0/src/tacl/modules/thought_block.py +65 -0
- tacl_llm-2.0.0/src/tacl/presets.py +80 -0
- tacl_llm-2.0.0/src/tacl/scripts/__init__.py +0 -0
- tacl_llm-2.0.0/src/tacl/scripts/build_vocab.py +108 -0
- tacl_llm-2.0.0/src/tacl/scripts/generate.py +67 -0
- tacl_llm-2.0.0/src/tacl/scripts/import_vocab.py +46 -0
- tacl_llm-2.0.0/src/tacl/scripts/smoke_check.py +303 -0
- tacl_llm-2.0.0/src/tacl/scripts/train.py +106 -0
- tacl_llm-2.0.0/src/tacl/serialization.py +227 -0
- tacl_llm-2.0.0/src/tacl/tokenization_tacl.py +223 -0
- tacl_llm-2.0.0/src/tacl/training/__init__.py +11 -0
- tacl_llm-2.0.0/src/tacl/training/callbacks.py +59 -0
- tacl_llm-2.0.0/src/tacl/training/chl.py +183 -0
- tacl_llm-2.0.0/src/tacl/training/logging_utils.py +99 -0
- tacl_llm-2.0.0/src/tacl/training/trainer.py +306 -0
- tacl_llm-2.0.0/src/tacl/utils/__init__.py +4 -0
- tacl_llm-2.0.0/src/tacl/utils/device.py +16 -0
- tacl_llm-2.0.0/src/tacl/utils/seed.py +11 -0
- tacl_llm-2.0.0/src/tacl_llm.egg-info/PKG-INFO +26 -0
- tacl_llm-2.0.0/src/tacl_llm.egg-info/SOURCES.txt +46 -0
- tacl_llm-2.0.0/src/tacl_llm.egg-info/dependency_links.txt +1 -0
- tacl_llm-2.0.0/src/tacl_llm.egg-info/entry_points.txt +4 -0
- tacl_llm-2.0.0/src/tacl_llm.egg-info/requires.txt +24 -0
- tacl_llm-2.0.0/src/tacl_llm.egg-info/top_level.txt +1 -0
tacl_llm-2.0.0/PKG-INFO
ADDED
|
@@ -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"
|
tacl_llm-2.0.0/README.md
ADDED
|
@@ -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
|
+
]
|
tacl_llm-2.0.0/setup.cfg
ADDED
|
@@ -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}
|