| import torch
|
|
|
| from taoTrain.data.async_loader import AsyncBatchIterator
|
| from taoTrain.data.sft_jsonl import SFTJSONLDataset
|
| from taoTrain.data.sft_utils import build_response_only_next_token_labels
|
|
|
|
|
| def test_response_only_labels_mask_target_token_not_input_token():
|
| input_ids = [10, 20, 30, 40, 2, 0]
|
| mask = [0, 0, 1, 1, 0, 0]
|
|
|
| labels = build_response_only_next_token_labels(input_ids, mask)
|
|
|
| assert labels == [-100, 30, 40, -100, -100, -100]
|
|
|
|
|
| def test_sft_dataset_direct_path_matches_response_only_helper():
|
| dataset = SFTJSONLDataset.__new__(SFTJSONLDataset)
|
| dataset.chunk_manager = None
|
| dataset._current_chunk_data = {
|
| "input_ids": [[10, 20, 30, 40, 2, 0]],
|
| "attention_mask": [[1, 1, 1, 1, 1, 0]],
|
| "mask": [[0, 0, 1, 1, 0, 0]],
|
| }
|
|
|
| sample = dataset[0]
|
|
|
| assert torch.equal(sample["labels"], torch.tensor([-100, 30, 40, -100, -100, -100]))
|
|
|
|
|
| class _OneChunkQueue:
|
| def __init__(self):
|
| self._chunk = {
|
| "input_ids": [[10, 20, 30, 40, 2, 0]],
|
| "attention_mask": [[1, 1, 1, 1, 1, 0]],
|
| "mask": [[0, 0, 1, 1, 0, 0]],
|
| }
|
| self._returned = False
|
| self._next_chunk_idx = 0
|
| self._chunk_order = [0]
|
| self._threads = [object()]
|
|
|
| def get_next_chunk(self, timeout=None):
|
| if self._returned:
|
| return None
|
| self._returned = True
|
| return self._chunk
|
|
|
| @property
|
| def is_exhausted(self):
|
| return self._returned
|
|
|
| def shutdown(self, wait=True):
|
| return None
|
|
|
| def __len__(self):
|
| return 1
|
|
|
|
|
| def test_async_sft_loader_matches_direct_dataset_labels():
|
| iterator = AsyncBatchIterator(
|
| tokenization_queue=_OneChunkQueue(),
|
| batch_size=1,
|
| device=torch.device("cpu"),
|
| drop_last=True,
|
| )
|
|
|
| batch = next(iter(iterator))
|
|
|
| assert torch.equal(batch["labels"], torch.tensor([[-100, 30, 40, -100, -100, -100]]))
|
|
|
