tacl-llm 2.0.0__py3-none-any.whl

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/__init__.py 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"
tacl/auto.py ADDED
@@ -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
tacl/data/__init__.py ADDED
@@ -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
+ ]
tacl/data/base.py ADDED
@@ -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}
@@ -0,0 +1,200 @@
1
+ from __future__ import annotations
2
+
3
+ import csv
4
+ import json
5
+ from pathlib import Path
6
+ from typing import Any, Optional, Union
7
+
8
+ import torch
9
+ from torch import Tensor
10
+
11
+ from .base import TACLDataset
12
+
13
+ DEFAULT_TACL_CHAT_TEMPLATE = "<bos>system\n{system}\n<eos>\n<bos>user\n{user}\n<eos>\n<bos>assistant\n{assistant}<eos>"
14
+ _TEMPLATE_END = "<eos>\n<bos>assistant\n"
15
+
16
+
17
+ class SFTDataset(TACLDataset):
18
+ """
19
+ SFT / QA dataset supporting two input formats:
20
+ 1. {"instruction": "...", "input": "...", "output": "..."}
21
+ 2. {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
22
+
23
+ Supports:
24
+ - Local .jsonl files
25
+ - Local .csv files (columns: instruction, input, output)
26
+ - Local .parquet files (columns: instruction, input, output)
27
+ - HF Datasets repo_id
28
+ - List of dicts
29
+ - prompt_template for chat formatting
30
+ - mask_prompt_loss to zero out CHL updates on prompt tokens
31
+ """
32
+
33
+ def __init__(
34
+ self,
35
+ source: Union[str, list[dict]],
36
+ tokenizer: Any,
37
+ max_length: int = 512,
38
+ prompt_template: str = DEFAULT_TACL_CHAT_TEMPLATE,
39
+ mask_prompt_loss: bool = True,
40
+ ):
41
+ self._prompt_template = prompt_template
42
+ self._mask_prompt_loss = mask_prompt_loss
43
+ super().__init__(source, tokenizer, max_length)
44
+
45
+ def _load(self, source: Union[str, list[dict]]) -> None:
46
+ if isinstance(source, list):
47
+ raw = source
48
+ elif isinstance(source, str) and Path(source).exists():
49
+ ext = Path(source).suffix.lower()
50
+ if ext == ".jsonl":
51
+ with open(source, "r", encoding="utf-8") as f:
52
+ raw = [json.loads(line) for line in f if line.strip()]
53
+ elif ext == ".csv":
54
+ raw = []
55
+ with open(source, "r", encoding="utf-8", newline="") as f:
56
+ reader = csv.DictReader(f)
57
+ for row in reader:
58
+ raw.append(row)
59
+ elif ext == ".parquet":
60
+ try:
61
+ import pandas as pd
62
+ df = pd.read_parquet(source)
63
+ raw = df.to_dict("records")
64
+ except ImportError:
65
+ raise ImportError("Install pandas and pyarrow for .parquet support: pip install pandas pyarrow")
66
+ else:
67
+ with open(source, "r", encoding="utf-8") as f:
68
+ raw = [json.loads(line) for line in f if line.strip()]
69
+ elif isinstance(source, str):
70
+ try:
71
+ from datasets import load_dataset
72
+ ds = load_dataset(source, split="train")
73
+ raw = [dict(ex) for ex in ds]
74
+ except Exception:
75
+ raise ValueError(f"Cannot load dataset from {source}")
76
+ else:
77
+ raise TypeError(f"Unsupported source type: {type(source)}")
78
+
79
+ for example in raw:
80
+ tokenized = self._tokenize_example(example)
81
+ if tokenized is not None:
82
+ self._samples.append(tokenized)
83
+
84
+ def _tokenize_example(self, example: dict) -> Optional[dict[str, Tensor]]:
85
+ if "messages" in example:
86
+ return self._tokenize_chat(example["messages"])
87
+ elif "instruction" in example:
88
+ return self._tokenize_instruction(example)
89
+ else:
90
+ return None
91
+
92
+ def _tokenize_instruction(self, example: dict) -> Optional[dict[str, Tensor]]:
93
+ instruction = example.get("instruction", "")
94
+ input_text = example.get("input", "")
95
+ output = example.get("output", "")
96
+
97
+ user_text = instruction
98
+ if input_text:
99
+ user_text = f"{instruction}\n{input_text}"
100
+
101
+ prompt = self._prompt_template.format(
102
+ system="", user=user_text, assistant=""
103
+ )
104
+ prompt_tail = _TEMPLATE_END
105
+ assistant_start = prompt.find(prompt_tail)
106
+ if assistant_start != -1:
107
+ prompt = prompt[:assistant_start + len(prompt_tail)]
108
+
109
+ eos_id = self._get_eos_id()
110
+ full_text = prompt + output
111
+ if eos_id is not None:
112
+ full_text = full_text + "<eos>"
113
+
114
+ ids = self._tokenizer.encode(full_text)
115
+ prompt_ids = self._tokenizer.encode(prompt)
116
+
117
+ if len(ids) > self._max_length:
118
+ ids = ids[:self._max_length]
119
+
120
+ labels = list(ids)
121
+ if self._mask_prompt_loss and len(prompt_ids) > 0:
122
+ for i in range(min(len(prompt_ids) - 1, len(labels))):
123
+ labels[i] = -100
124
+
125
+ input_tensor = torch.tensor(ids, dtype=torch.long)
126
+ labels_tensor = torch.tensor(labels, dtype=torch.long)
127
+
128
+ if len(input_tensor) < 2:
129
+ return None
130
+
131
+ return {"input_ids": input_tensor, "labels": labels_tensor}
132
+
133
+ def _tokenize_chat(self, messages: list[dict]) -> Optional[dict[str, Tensor]]:
134
+ user_text = ""
135
+ assistant_text = ""
136
+ for msg in messages:
137
+ if msg["role"] == "user":
138
+ user_text = msg["content"]
139
+ elif msg["role"] == "assistant":
140
+ assistant_text = msg["content"]
141
+
142
+ prompt = self._prompt_template.format(
143
+ system="", user=user_text, assistant=""
144
+ )
145
+ prompt_tail = _TEMPLATE_END
146
+ assistant_start = prompt.find(prompt_tail)
147
+ if assistant_start != -1:
148
+ prompt = prompt[:assistant_start + len(prompt_tail)]
149
+
150
+ eos_id = self._get_eos_id()
151
+ full_text = prompt + assistant_text
152
+ if eos_id is not None:
153
+ full_text = full_text + "<eos>"
154
+
155
+ ids = self._tokenizer.encode(full_text)
156
+ prompt_ids = self._tokenizer.encode(prompt)
157
+
158
+ if len(ids) > self._max_length:
159
+ ids = ids[:self._max_length]
160
+
161
+ labels = list(ids)
162
+ if self._mask_prompt_loss and len(prompt_ids) > 0:
163
+ for i in range(min(len(prompt_ids), len(labels))):
164
+ labels[i] = -100
165
+
166
+ input_tensor = torch.tensor(ids, dtype=torch.long)
167
+ labels_tensor = torch.tensor(labels, dtype=torch.long)
168
+
169
+ if len(input_tensor) < 2:
170
+ return None
171
+
172
+ return {"input_ids": input_tensor, "labels": labels_tensor}
173
+
174
+
175
+ def sft_collate_fn(batch: list[dict[str, Tensor]]) -> dict[str, Tensor]:
176
+ max_len = max(item["input_ids"].size(0) for item in batch)
177
+ vocab_size = 0
178
+
179
+ input_ids_list = []
180
+ labels_list = []
181
+ attention_mask_list = []
182
+
183
+ for item in batch:
184
+ ids = item["input_ids"]
185
+ labs = item["labels"]
186
+ pad_len = max_len - ids.size(0)
187
+
188
+ padded_ids = torch.cat([ids, torch.zeros(pad_len, dtype=torch.long)])
189
+ padded_labs = torch.cat([labs, torch.full((pad_len,), -100, dtype=torch.long)])
190
+ mask = torch.cat([torch.ones(ids.size(0)), torch.zeros(pad_len)])
191
+
192
+ input_ids_list.append(padded_ids)
193
+ labels_list.append(padded_labs)
194
+ attention_mask_list.append(mask)
195
+
196
+ return {
197
+ "input_ids": torch.stack(input_ids_list),
198
+ "labels": torch.stack(labels_list),
199
+ "attention_mask": torch.stack(attention_mask_list),
200
+ }
tacl/data/streaming.py ADDED
@@ -0,0 +1,52 @@
1
+ from __future__ import annotations
2
+
3
+ import logging
4
+ from pathlib import Path
5
+ from typing import Any, Iterator, Optional
6
+
7
+ import torch
8
+ from torch.utils.data import IterableDataset
9
+
10
+ logger = logging.getLogger(__name__)
11
+
12
+
13
+ class StreamingTextDataset(IterableDataset):
14
+ """
15
+ Memory-mapped streaming dataset for pre-tokenized binary data.
16
+ Expects a numpy .npy or .memmap file of uint16 token IDs.
17
+ """
18
+
19
+ def __init__(
20
+ self,
21
+ data_path: str,
22
+ seq_len: int = 1024,
23
+ tokenizer: Any = None,
24
+ ):
25
+ self._data_path = data_path
26
+ self._seq_len = seq_len
27
+ self._tokenizer = tokenizer
28
+ self._data = None
29
+
30
+ def _load_data(self):
31
+ import numpy as np
32
+ path = Path(self._data_path)
33
+ if path.suffix == ".npy":
34
+ self._data = np.load(str(path)).astype(np.int64)
35
+ else:
36
+ self._data = np.memmap(str(path), dtype=np.uint16, mode="r")
37
+
38
+ def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
39
+ if self._data is None:
40
+ self._load_data()
41
+
42
+ for i in range(0, len(self._data) - self._seq_len, self._seq_len):
43
+ chunk = torch.from_numpy(self._data[i:i + self._seq_len].astype(int))
44
+ yield {
45
+ "input_ids": chunk.long(),
46
+ "labels": chunk.long(),
47
+ }
48
+
49
+ def __len__(self) -> int:
50
+ if self._data is None:
51
+ self._load_data()
52
+ return max(0, (len(self._data) - 1) // self._seq_len)
tacl/exceptions.py ADDED
@@ -0,0 +1,9 @@
1
+ import warnings
2
+
3
+
4
+ class TACLCaveatWarning(UserWarning):
5
+ pass
6
+
7
+
8
+ def warn_caveat(message: str) -> None:
9
+ warnings.warn(f"[TACL caveat] {message}", TACLCaveatWarning, stacklevel=2)