s444415/donut
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Custom wrapper

Browse Source
This commit is contained in:
zzombely 2023-03-14 22:01:50 +00:00
parent 30369a7885
commit ef0faa1f8c
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
train_stream.py
View File

@ -19,7 +19,7 @@ from torchdata.datapipes.iter import IterableWrapper
import json
class TestIterator(IterableWrapper):
class CustomWrapperIterator(IterableWrapper):
def __init__(self, iterable, deepcopy=True, total_len=None):
super().__init__(iterable, deepcopy)
self.total_len = total_len
@ -181,8 +181,8 @@ def main(config, hug_token):
# train_dataset = train_dataset.with_format('torch')
# val_dataset = val_dataset.with_format('torch')
train_dataset = TestIterator(train_dataset, total_len=train_length)
val_dataset = TestIterator(val_dataset, total_len=val_length)
train_dataset = CustomWrapperIterator(train_dataset, total_len=train_length)
val_dataset = CustomWrapperIterator(val_dataset, total_len=val_length)
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]