abacaj / fine-tune-mistral Goto Github PK

class DummyTokenizer:
    pad_token_id = 0

@dataclass
class DataCollatorForSupervisedDataset(object):
    """Collate examples for supervised fine-tuning."""

    tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids, labels = tuple(
            [instance[key] for instance in instances] for key in ("input_ids", "labels")
        )
        
        # BEGIN: added line to return torch.tensor
        input_ids = [torch.tensor(x) for x in input_ids]
        labels = [torch.tensor(x) for x in labels]        
        # END

        input_ids = torch.nn.utils.rnn.pad_sequence(
            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
        )
        labels = torch.nn.utils.rnn.pad_sequence(
            labels, batch_first=True, padding_value=-100
        )

        return dict(
            input_ids=input_ids,
            labels=labels,
            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
        )

ds = [(torch.ones(x)*x).long() for x in np.random.permutation(np.arange(1,101))]
ds = [{"input_ids":x, "labels":x} for x in ds]
ds = datasets.Dataset.from_list(ds)
lengths = np.array([len(x['input_ids']) for x in ds])

train_sampler = MultipackDistributedBatchSampler(
    batch_max_length=4*128,
    lengths=lengths,
    num_replicas=1,
    rank=0,
    seed=42,
)

tokenizer = DummyTokenizer()
collator = DataCollatorForSupervisedDataset(tokenizer)
train_loader = DataLoader(
    ds,
    pin_memory=False,
    collate_fn=collator,
    batch_sampler=train_sampler,
)

for b in train_loader:
    print((b['input_ids'] != tokenizer.pad_token_id).view(-1).float().mean())
    print(train_loader.batch_sampler.efficiency())


tensor(0.5262)
0.8966619318181818
tensor(0.6364)
0.8966619318181818
tensor(0.4837)
0.8966619318181818
tensor(0.4582)
0.8966619318181818
tensor(0.6306)
0.8966619318181818
tensor(0.7002)
0.8966619318181818
tensor(0.5488)
0.8966619318181818
tensor(0.5200)
0.8966619318181818
tensor(0.4594)
0.8966619318181818
tensor(0.8535)
0.8966619318181818
tensor(0.5618)
0.8966619318181818