Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
bert_dataset.py 2.05 KB
Newer Older
yuguo960516's avatar
bloom  
yuguo960516 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66 
from libai.config import LazyCall
from omegaconf import OmegaConf
from libai.data import build_nlp_test_loader, build_nlp_train_val_test_loader
from libai.data.datasets import BertDataset
from libai.data.data_utils import get_indexed_dataset

from libai.tokenizer import BertTokenizer


tokenization = OmegaConf.create()

tokenization.tokenizer = LazyCall(BertTokenizer)(
    vocab_file="bert-base-chinese-vocab.txt",
    do_lower_case=True,
    do_chinese_wwm=True,
)
tokenization.append_eod = False
tokenization.make_vocab_size_divisible_by = 128

dataloader = OmegaConf.create()

dataloader.train = LazyCall(build_nlp_train_val_test_loader)(
    dataset=[
        LazyCall(BertDataset)(
            name="bert",
            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
            indexed_dataset=LazyCall(get_indexed_dataset)(
                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
                data_impl="mmap",
                skip_warmup=False,
            ),
            max_seq_length=512,
            mask_lm_prob=0.15,
            short_seq_prob=0.1,
            binary_head=True,
            seed=1234,
            masking_style="bert-cn-wwm",
        ),
    ],
    train_val_test_num_samples=None,  # a hint for deferred assignment
    splits=[[949.0, 50.0, 1.0]],
    weights=[1.0],
    num_workers=4,
)

dataloader.test = [
    LazyCall(build_nlp_test_loader)(
        dataset=LazyCall(BertDataset)(
            name="bert",
            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
            indexed_dataset=LazyCall(get_indexed_dataset)(
                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
                data_impl="mmap",
                skip_warmup=False,
            ),
            max_num_samples=10,
            max_seq_length=512,
            mask_lm_prob=0.15,
            short_seq_prob=0.1,
            binary_head=True,
            seed=1234,
            masking_style="bert-cn-wwm",
        ),
        test_batch_size=4,
    )
]