tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

set labels and fix datasets block

Browse Source
This commit is contained in:
Wing Lian
2024-12-13 13:04:24 -05:00
parent d000851eeb
commit 9eaae5925a
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/prompt_strategies/pretrain.py
View File

@@ -41,6 +41,7 @@ class PretrainTokenizationStrategy(PromptTokenizingStrategy):
seq + [self.tokenizer.eos_token_id] for seq in res["input_ids"] seq + [self.tokenizer.eos_token_id] for seq in res["input_ids"]
] ]
res["attention_mask"] = [seq + [1] for seq in res["attention_mask"]] res["attention_mask"] = [seq + [1] for seq in res["attention_mask"]]
res["labels"] = res["input_ids"].copy()
return res return res
@@ -52,7 +53,7 @@ def load(tokenizer, cfg):
if cfg.pretraining_dataset: if cfg.pretraining_dataset:
cfg_ds = cfg.pretraining_dataset cfg_ds = cfg.pretraining_dataset
else: else:
cfg_ds = cfg.dataset cfg_ds = cfg.datasets
strat = PretrainTokenizationStrategy( strat = PretrainTokenizationStrategy(
PretrainTokenizer(), PretrainTokenizer(),
tokenizer, tokenizer,