/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
examples/tutorial/sequence_parallel/data/dummy_dataloader.py
52 строки
1 KB
Hongxin Liu
[misc] update pre-commit and run all files (#4752)
19 сен 2023, 09:20
Не верифицирован
19 сен 2023, 09:20
079bf3c
Код
Авторство
О чём код?
import torch class DummyDataloader: def __init__(self, batch_size, vocab_size, seq_length): self.batch_size = batch_size self.vocab_size = vocab_size self.seq_length = seq_length self.step = 0 def generate(self): tokens = torch.randint( low=0, high=self.vocab_size, size=( self.batch_size, self.seq_length, ), ) types = torch.randint( low=0, high=3, size=( self.batch_size, self.seq_length, ), ) sentence_order = torch.randint(low=0, high=2, size=(self.batch_size,)) loss_mask = torch.randint( low=0, high=2, size=( self.batch_size, self.seq_length, ), ) lm_labels = torch.randint(low=0, high=self.vocab_size, size=(self.batch_size, self.seq_length)) padding_mask = torch.randint(low=0, high=2, size=(self.batch_size, self.seq_length)) return dict( text=tokens, types=types, is_random=sentence_order, loss_mask=loss_mask, labels=lm_labels, padding_mask=padding_mask, ) def __iter__(self): return self def __next__(self): return self.generate()