slovnet

1

2
from os import getenv, environ
3
from os.path import exists, join
4
from itertools import chain, islice as head
5
from random import seed, sample, randint, uniform
6
from subprocess import run
7

8
from tqdm.notebook import tqdm as log_progress
9

10
import torch
11
from torch import optim
12

13
from naeval.syntax.datasets import load_dataset
14

15
from slovnet.s3 import S3
16
from slovnet.io import (
17
    format_jl,
18
    parse_jl,
19

20
    load_gz_lines,
21
    dump_gz_lines
22
)
23
from slovnet.board import (
24
    TensorBoard,
25
    LogBoard,
26
    MultiBoard
27
)
28
from slovnet.const import (
29
    TRAIN, TEST,
30
    PAD, CUDA0,
31
)
32

33
from slovnet.model.bert import (
34
    RuBERTConfig,
35
    BERTEmbedding,
36
    BERTEncoder,
37
    BERTSyntaxHead,
38
    BERTSyntaxRel,
39
    BERTSyntax
40
)
41
from slovnet.markup import SyntaxMarkup
42
from slovnet.vocab import BERTVocab, Vocab
43
from slovnet.encoders.bert import BERTSyntaxTrainEncoder
44
from slovnet.loss import masked_flatten_cross_entropy
45
from slovnet.score import (
46
    SyntaxScoreMeter,
47
    score_syntax_batch
48
)
49
from slovnet.mask import (
50
    Masked,
51
    split_masked,
52
    pad_masked
53
)
54

55

56
DATA_DIR = 'data'
57
MODEL_DIR = 'model'
58
BERT_DIR = 'bert'
59
RAW_DIR = join(DATA_DIR, 'raw')
60

61
NEWS = join(DATA_DIR, 'news.jl.gz')
62
FICTION = join(DATA_DIR, 'fiction.jl.gz')
63
GRAMRU_DIR = join(RAW_DIR, 'GramEval2020-master')
64
GRAMRU_FILES = {
65
    NEWS: [
66
        'dataOpenTest/GramEval2020-RuEval2017-Lenta-news-dev.conllu',
67
        'dataTrain/MorphoRuEval2017-Lenta-train.conllu',
68
    ],
69
    FICTION: [
70
        'dataOpenTest/GramEval2020-SynTagRus-dev.conllu',
71
        'dataTrain/GramEval2020-SynTagRus-train-v2.conllu',
72
        'dataTrain/MorphoRuEval2017-JZ-gold.conllu'
73
    ],
74
}
75

76
S3_DIR = '04_bert_syntax'
77
S3_NEWS = join(S3_DIR, NEWS)
78
S3_FICTION = join(S3_DIR, FICTION)
79

80
VOCAB = 'vocab.txt'
81
EMB = 'emb.pt'
82
ENCODER = 'encoder.pt'
83
HEAD = 'head.pt'
84
REL = 'rel.pt'
85

86
BERT_VOCAB = join(BERT_DIR, VOCAB)
87
BERT_EMB = join(BERT_DIR, EMB)
88
BERT_ENCODER = join(BERT_DIR, ENCODER)
89

90
S3_RUBERT_DIR = '01_bert_news/rubert'
91
S3_MLM_DIR = '01_bert_news/model'
92
S3_BERT_VOCAB = join(S3_RUBERT_DIR, VOCAB)
93
S3_BERT_EMB = join(S3_MLM_DIR, EMB)
94
S3_BERT_ENCODER = join(S3_MLM_DIR, ENCODER)
95

96
RELS_VOCAB = join(MODEL_DIR, 'rels_vocab.txt')
97
MODEL_ENCODER = join(MODEL_DIR, ENCODER)
98
MODEL_HEAD = join(MODEL_DIR, HEAD)
99
MODEL_REL = join(MODEL_DIR, REL)
100

101
S3_RELS_VOCAB = join(S3_DIR, RELS_VOCAB)
102
S3_MODEL_ENCODER = join(S3_DIR, MODEL_ENCODER)
103
S3_MODEL_HEAD = join(S3_DIR, MODEL_HEAD)
104
S3_MODEL_REL = join(S3_DIR, MODEL_REL)
105

106
BOARD_NAME = getenv('board_name', '04_bert_syntax_01')
107
RUNS_DIR = 'runs'
108

109
TRAIN_BOARD = '01_train'
110
TEST_BOARD = '02_test'
111

112
SEED = int(getenv('seed', 50))
113
DEVICE = getenv('device', CUDA0)
114
BERT_LR = float(getenv('bert_lr', 0.000058))
115
LR = float(getenv('lr', 0.00012))
116
LR_GAMMA = float(getenv('lr_gamma', 0.29))
117
EPOCHS = int(getenv('epochs', 2))
118

119

120
def process_batch(model, criterion, batch):
121
    input, target = batch
122

123
    pred = model(
124
        input.word_id, input.word_mask, input.pad_mask,
125
        target.mask, target.head_id
126
    )
127

128
    loss = (
129
        criterion(pred.head_id, target.head_id, target.mask)
130
        + criterion(pred.rel_id, target.rel_id, target.mask)
131
    )
132

133
    pred.head_id = model.head.decode(pred.head_id, target.mask)
134
    pred.rel_id = model.rel.decode(pred.rel_id, target.mask)
135

136
    return batch.processed(loss, pred)
137