CSS-LM

1
import argparse
2
import logging
3
import random
4
import numpy as np
5
import os
6
import json
7
import math
8

9
import torch
10
from transformers import RobertaTokenizer, RobertaForMaskedLM, RobertaForSequenceClassification
11
#from transformers.modeling_roberta import RobertaForMaskedLMDomainTask
12
from transformers.modeling_roberta_updateRep_self import RobertaForMaskedLMDomainTask
13
from tqdm import tqdm, trange
14
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
15
from torch.utils.data.distributed import DistributedSampler
16
from transformers.file_utils import PYTORCH_PRETRAINED_BERT_CACHE
17
from transformers.optimization import AdamW, get_linear_schedule_with_warmup
18

19

20

21
logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
22
                    datefmt = '%m/%d/%Y %H:%M:%S',
23
                    level = logging.INFO)
24
logger = logging.getLogger(__name__)
25

26
def accuracy(out, labels):
27
    outputs = np.argmax(out, axis=1)
28
    return np.sum(outputs == labels), outputs
29

30
class InputFeatures(object):
31
    """A single set of features of data."""
32

33
    def __init__(self, input_ids=None, attention_mask=None, segment_ids=None, label_id=None):
34
        self.input_ids = input_ids
35
        self.attention_mask = attention_mask
36
        self.segment_ids = segment_ids
37
        self.label_id = label_id
38

39
class InputExample(object):
40
    """A single training/test example for simple sequence classification."""
41
    def __init__(self, guid, sentence, aspect, sentiment=None):
42
        """Constructs a InputExample.
43

44
        Args:
45
            guid: Unique id for the example.
46
            text_a: string. The untokenized text of the first sequence. For single
47
            sequence tasks, only this sequence must be specified.
48
            text_b: (Optional) string. The untokenized text of the second sequence.
49
            Only must be specified for sequence pair tasks.
50
            label: (Optional) string. The label of the example. This should be
51
            specified for train and dev examples, but not for test examples.
52
        """
53
        self.guid = guid
54
        self.sentence = sentence
55
        self.aspect = aspect
56
        self.sentiment = sentiment
57

58

59
class DataProcessor(object):
60
    """Base class for data converters for sequence classification data sets."""
61

62
    def get_train_examples(self, data_dir):
63
        """Gets a collection of `InputExample`s for the train set."""
64
        raise NotImplementedError()
65

66
    def get_dev_examples(self, data_dir):
67
        """Gets a collection of `InputExample`s for the dev set."""
68
        raise NotImplementedError()
69

70
    def get_labels(self):
71
        """Gets the list of labels for this data set."""
72
        raise NotImplementedError()
73

74
    @classmethod
75
    def _read_json(cls, input_file):
76
        with open(input_file, "r", encoding='utf-8') as f:
77
            return json.loads(f.read())
78

79

80
class Processor_1(DataProcessor):
81
    """Processor for the CoLA data set (GLUE version)."""
82

83
    def get_train_examples(self, data_dir):
84
        """See base class."""
85
        examples = self._create_examples(
86
            self._read_json(os.path.join(data_dir, "train.json")), "train")
87
        aspect = set([x.aspect for x in examples])
88
        sentiment = set([x.sentiment for x in examples])
89
        return examples, list(aspect), list(sentiment)
90

91
    def get_dev_examples(self, data_dir):
92
        """See base class."""
93
        examples = self._create_examples(
94
            self._read_json(os.path.join(data_dir, "dev.json")), "dev")
95
        aspect = set([x.aspect for x in examples])
96
        sentiment = set([x.sentiment for x in examples])
97
        return examples, list(aspect), list(sentiment)
98

99
    def get_test_examples(self, data_dir):
100
        """See base class."""
101
        examples = self._create_examples(
102
            self._read_json(os.path.join(data_dir, "test.json")), "test")
103
        aspect = set([x.aspect for x in examples])
104
        sentiment = set([x.sentiment for x in examples])
105
        return examples, list(aspect), list(sentiment)
106

107
    def get_labels(self):
108
        """Useless"""
109
        return ["0", "1"]
110

111
    def _create_examples(self, lines, set_type):
112
        """Creates examples for the training and dev sets."""
113
        examples = []
114
        for (i, line) in enumerate(lines):
115
            guid = "%s-%s" % (set_type, i)
116

117
            sentence = line["sentence"]
118
            aspect = line["aspect"]
119
            sentiment = line["sentiment"]
120

121
            examples.append(
122
                InputExample(guid=guid, sentence=sentence, aspect=aspect, sentiment=sentiment))
123
        return examples
124

125
def convert_examples_to_features(examples, aspect_list, sentiment_list, max_seq_length, tokenizer, task_n):
126

127
    """Loads a data file into a list of `InputBatch`s."""
128

129
    #Task_1: sentence --> aspect
130
    #Task_2: aspect+sentence --> sentiment
131
    if task_n == 1:
132
        label_list = sorted(aspect_list)
133
    elif task_n == 2:
134
        label_list = sorted(sentiment_list)
135
    else:
136
        print("Wrong task")
137
    '''
138
    for w in label_list:
139
        print(w,tokenizer.encode(w))
140
    exit()
141
    '''
142
    label_map = {label : i for i, label in enumerate(label_list)}
143
    print("=======")
144
    print(label_map)
145
    print("=======")
146

147

148
    features = []
149
    for (ex_index, example) in enumerate(examples):
150

151
        #Add new special tokens
152
        '''
153
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
154
model = GPT2Model.from_pretrained('gpt2')
155
        special_tokens_dict = {'cls_token': '<CLS>'}
156
        num_added_toks = tokenizer.add_special_tokens(special_tokens_dict)
157
        print('We have added', num_added_toks, 'tokens')
158
        model.resize_token_embeddings(len(tokenizer))
159
        '''
160

161
        '''
162
        print(tokenizer.all_special_tokens)
163
        print(tokenizer.encode(tokenizer.all_special_tokens))
164
        #['[PAD]', '[SEP]', '[CLS]', '[MASK]', '[UNK]']
165
        #[ 0, 102, 101, 103, 100]
166
        '''
167

168

169
        # The convention in BERT is:
170
        # (a) For sequence pairs:
171
        #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
172
        #  type_ids: 0   0  0    0    0     0       0 0    1  1  1  1   1 1
173
        # (b) For single sequences:
174
        #  tokens:   [CLS] the dog is hairy . [SEP]
175
        #  type_ids: 0   0   0   0  0     0 0
176
        #
177
        # Where "type_ids" are used to indicate whether this is the first
178
        # sequence or the second sequence. The embedding vectors for `type=0` and
179
        # `type=1` were learned during pre-training and are added to the wordpiece
180
        # embedding vector (and position vector). This is not *strictly* necessary
181
        # since the [SEP] token unambigiously separates the sequences, but it makes
182
        # it easier for the model to learn the concept of sequences.
183
        #
184
        # For classification tasks, the first vector (corresponding to [CLS]) is
185
        # used as as the "sentence vector". Note that this only makes sense because
186
        # the entire model is fine-tuned.
187

188
        ###
189
        #Already add [CLS] and [SEP]
190
        #101, 102
191
        input_ids = tokenizer.encode(example.sentence,add_special_tokens=True)
192
        if len(input_ids) > max_seq_length:
193
            input_ids = input_ids[:max_seq_length-1]+[2]
194
        segment_ids = [0] * len(input_ids)
195

196

197
        '''
198
        if task_n==2:
199
            #"[SEP]"
200
            input_ids += input_ids + [102]
201
            #sentiment: word (Next sentence)
202
            #segment_ids += [1] * (len(tokens_b) + 1)
203
        '''
204

205
        # The “Attention Mask” is simply an array of 1s and 0s indicating which tokens are padding and which aren’t (including special tokens)
206

207
        # The mask has 1 for real tokens and 0 for padding tokens. Only real
208
        # tokens are attended to.
209
        attention_mask = [1] * len(input_ids)
210

211
        # Zero-pad up to the sequence length.
212
        padding = [0] * (max_seq_length - len(input_ids))
213
        #<pad>:1
214
        padding_id = [1] * (max_seq_length - len(input_ids))
215
        input_ids += padding_id
216
        attention_mask += padding
217
        segment_ids += padding
218

219
        try:
220
            assert len(input_ids) == max_seq_length
221
            assert len(attention_mask) == max_seq_length
222
            assert len(segment_ids) == max_seq_length
223
        except:
224
            continue
225

226
        if task_n == 1:
227
            label_id = label_map[example.aspect]
228
        elif task_n == 2:
229
            label_id = label_map[example.sentiment]
230
        else:
231
            print("Wrong task")
232

233

234
        if task_n == 1:
235
            features.append(
236
                    InputFeatures(input_ids=input_ids,
237
                                  attention_mask=attention_mask,
238
                                  segment_ids=None,
239
                                  label_id=label_id))
240
        elif task_n == 2:
241
            features.append(
242
                    InputFeatures(input_ids=input_ids,
243
                                  attention_mask=attention_mask,
244
                                  segment_ids=segment_ids,
245
                                  label_id=label_id))
246
        else:
247
            print("Wrong in convert_examples")
248

249

250
    return features
251

252

253
def main():
254
    parser = argparse.ArgumentParser()
255
    ## Required parameters
256
    ###############
257
    parser.add_argument("--data_dir",
258
                        default=None,
259
                        type=str,
260
                        required=True,
261
                        help="The input data dir. Should contain the .tsv files (or other data files) for the task.")
262
    parser.add_argument("--output_dir",
263
                        default=None,
264
                        type=str,
265
                        required=True,
266
                        help="The output directory where the model predictions and checkpoints will be written.")
267
    parser.add_argument("--pretrain_model",
268
                        default='bert-case-uncased',
269
                        type=str,
270
                        required=True,
271
                        help="Pre-trained model")
272
    parser.add_argument("--num_labels_task",
273
                        default=None, type=int,
274
                        required=True,
275
                        help="num_labels_task")
276
    parser.add_argument("--max_seq_length",
277
                        default=128,
278
                        type=int,
279
                        help="The maximum total input sequence length after WordPiece tokenization. \n"
280
                             "Sequences longer than this will be truncated, and sequences shorter \n"
281
                             "than this will be padded.")
282
    parser.add_argument("--do_train",
283
                        default=False,
284
                        action='store_true',
285
                        help="Whether to run training.")
286
    parser.add_argument("--do_eval",
287
                        default=False,
288
                        action='store_true',
289
                        help="Whether to run eval on the dev set.")
290
    parser.add_argument("--do_lower_case",
291
                        default=False,
292
                        action='store_true',
293
                        help="Set this flag if you are using an uncased model.")
294
    parser.add_argument("--eval_batch_size",
295
                        default=32,
296
                        type=int,
297
                        help="Total batch size for training.")
298
    parser.add_argument("--learning_rate",
299
                        default=5e-5,
300
                        type=float,
301
                        help="The initial learning rate for Adam.")
302
    parser.add_argument("--num_train_epochs",
303
                        default=3.0,
304
                        type=float,
305
                        help="Total number of training epochs to perform.")
306
    parser.add_argument("--warmup_proportion",
307
                        default=0.1,
308
                        type=float,
309
                        help="Proportion of training to perform linear learning rate warmup for. "
310
                             "E.g., 0.1 = 10%% of training.")
311
    parser.add_argument("--no_cuda",
312
                        default=False,
313
                        action='store_true',
314
                        help="Whether not to use CUDA when available")
315
    parser.add_argument("--local_rank",
316
                        type=int,
317
                        default=-1,
318
                        help="local_rank for distributed training on gpus")
319
    parser.add_argument('--seed',
320
                        type=int,
321
                        default=42,
322
                        help="random seed for initialization")
323
    parser.add_argument('--gradient_accumulation_steps',
324
                        type=int,
325
                        default=1,
326
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
327
    parser.add_argument('--fp16',
328
                        default=False,
329
                        action='store_true',
330
                        help="Whether to use 16-bit float precision instead of 32-bit")
331
    parser.add_argument('--loss_scale',
332
                        type=float, default=0,
333
                        help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
334
                             "0 (default value): dynamic loss scaling.\n"
335
                             "Positive power of 2: static loss scaling value.\n")
336
    parser.add_argument("--weight_decay",
337
                        default=0.0,
338
                        type=float,
339
                        help="Weight decay if we apply some.")
340
    parser.add_argument("--adam_epsilon",
341
                        default=1e-8,
342
                        type=float,
343
                        help="Epsilon for Adam optimizer.")
344
    parser.add_argument("--max_grad_norm",
345
                        default=1.0,
346
                        type=float,
347
                        help="Max gradient norm.")
348
    parser.add_argument('--fp16_opt_level',
349
                        type=str,
350
                        default='O1',
351
                        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
352
                             "See details at https://nvidia.github.io/apex/amp.html")
353
    parser.add_argument("--task",
354
                        default=2,
355
                        type=int,
356
                        required=True,
357
                        help="Choose Task")
358
    parser.add_argument("--choose_eval_test_both",
359
                        default=2,
360
                        type=int,
361
                        help="choose test dev both")
362
    ###############
363

364
    args = parser.parse_args()
365
    #print(args.do_train, args.do_eval)
366
    #exit()
367

368

369
    processors = Processor_1
370

371
    num_labels = args.num_labels_task
372

373
    if args.local_rank == -1 or args.no_cuda:
374
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
375
        n_gpu = torch.cuda.device_count()
376
        print(n_gpu)
377
        print(device)
378
    else:
379
        torch.cuda.set_device(args.local_rank)
380
        device = torch.device("cuda", args.local_rank)
381
        n_gpu = 1
382
        # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
383
        torch.distributed.init_process_group(backend='nccl')
384
    logger.info("device: {}, n_gpu: {}, distributed training: {}, 16-bits training: {}".format(
385
        device, n_gpu, bool(args.local_rank != -1), args.fp16))
386

387

388

389
    if args.gradient_accumulation_steps < 1:
390
        raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
391
                            args.gradient_accumulation_steps))
392

393
    #args.train_batch_size = int(args.train_batch_size / args.gradient_accumulation_steps)
394

395
    random.seed(args.seed)
396
    np.random.seed(args.seed)
397
    torch.manual_seed(args.seed)
398
    if n_gpu > 0:
399
        torch.cuda.manual_seed_all(args.seed)
400

401
    if not args.do_eval:
402
        raise ValueError("At least one of `do_train` or `do_eval` must be True.")
403

404
    '''
405
    if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train:
406
        raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
407
    '''
408
    os.makedirs(args.output_dir, exist_ok=True)
409

410

411

412
    tokenizer = RobertaTokenizer.from_pretrained(args.pretrain_model)
413

414

415
    train_examples = None
416
    num_train_steps = None
417
    aspect_list = None
418
    sentiment_list = None
419
    processor = processors()
420
    num_labels = num_labels
421
    #train_examples, aspect_list, sentiment_list = processor.get_train_examples(args.data_dir)
422

423
    filenames = os.listdir(args.output_dir)
424
    filenames = [x for x in filenames if "pytorch_model.bin_" in x]
425
    print(filenames)
426

427
    file_mark = []
428
    model_performace_dev = dict()
429
    model_performace_test = dict()
430
    for x in filenames:
431
        ###
432
        #test
433
        if args.choose_eval_test_both==0:
434
            file_mark.append([x, True])
435
        #eval
436
        elif args.choose_eval_test_both==1:
437
            file_mark.append([x, False])
438
        else:
439
            file_mark.append([x, True])
440
            file_mark.append([x, False])
441
        ###
442
        #file_mark.append([x, True])
443
        #file_mark.append([x, False])
444

445
    ####
446
    ####
447
    train_examples, aspect_list, sentiment_list = processor.get_test_examples(args.data_dir)
448
    test_examples, _, _ = processor.get_test_examples(args.data_dir)
449
    eval_examples, _, _ = processor.get_dev_examples(args.data_dir)
450
    if args.task == 1:
451
        num_labels = len(aspect_list)
452
    elif args.task == 2:
453
        num_labels = len(sentiment_list)
454
    else:
455
        print("What's task?")
456
        exit()
457
    test = convert_examples_to_features(
458
        test_examples, aspect_list, sentiment_list, args.max_seq_length, tokenizer, args.task)
459

460
    dev = convert_examples_to_features(
461
        eval_examples, aspect_list, sentiment_list, args.max_seq_length, tokenizer, args.task)
462
    ###
463

464

465
    for x, mark in file_mark:
466
        #mark: eval-True; test-False
467
        #choose_eval_test_both: eval-0, test-1, both-2
468
        print(x, mark)
469
        output_model_file = os.path.join(args.output_dir, x)
470

471
        #model = RobertaForSequenceClassification.from_pretrained(args.pretrain_model, num_labels=num_labels, output_hidden_states=False, output_attentions=False, return_dict=True)
472
        model = RobertaForMaskedLMDomainTask.from_pretrained(args.pretrain_model, output_hidden_states=False, output_attentions=False, return_dict=True, num_labels=args.num_labels_task)
473
        model.load_state_dict(torch.load(output_model_file), strict=False)
474
        #strict False: ignore non-matching keys
475
        model.to(device)
476

477
        #######################################
478
        param_optimizer = list(model.named_parameters())
479
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
480
        #no_decay = ['bias', 'LayerNorm.weight']
481
        no_grad = ['bert.encoder.layer.11.output.dense_ent', 'bert.encoder.layer.11.output.LayerNorm_ent']
482
        param_optimizer = [(n, p) for n, p in param_optimizer if not any(nd in n for nd in no_grad)]
483
        optimizer_grouped_parameters = [
484
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
485
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
486
            ]
487
        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
488
        #scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=int(t_total*0.1), num_training_steps=t_total)
489
        if args.fp16:
490
            try:
491
                from apex import amp
492
            except ImportError:
493
                raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
494
                exit()
495

496
            model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)
497

498

499
        # multi-gpu training (should be after apex fp16 initialization)
500
        if n_gpu > 1:
501
            model = torch.nn.DataParallel(model)
502

503
        # Distributed training (should be after apex fp16 initialization)
504
        if args.local_rank != -1:
505
            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
506
                                                              output_device=args.local_rank,
507
                                                              find_unused_parameters=True)
508
        #######################################
509

510

511
        #param_optimizer = [para[0] for para in model.named_parameters()]
512
        #param_optimizer = [para for para in model.named_parameters()][-2]
513
        #print(param_optimizer)
514

515
        if mark:
516
            eval_features = dev
517
        else:
518
            eval_features = test
519

520
        logger.info("***** Running evaluation *****")
521
        logger.info("  Num examples = %d", len(eval_examples))
522
        logger.info("  Batch size = %d", args.eval_batch_size)
523

524

525
        all_input_ids = torch.tensor([f.input_ids for f in eval_features], dtype=torch.long)
526
        all_attention_mask = torch.tensor([f.attention_mask for f in eval_features], dtype=torch.long)
527
        if args.task == 1:
528
            print("Excuting the task 1")
529
        elif args.task == 2:
530
            all_segment_ids = torch.tensor([f.segment_ids for f in eval_features], dtype=torch.long)
531
        else:
532
            print("Wrong here2")
533

534
        all_label_ids = torch.tensor([f.label_id for f in eval_features], dtype=torch.long)
535

536
        if args.task == 1:
537
            eval_data = TensorDataset(all_input_ids, all_attention_mask, all_label_ids)
538
        elif args.task == 2:
539
            eval_data = TensorDataset(all_input_ids, all_attention_mask, all_segment_ids, all_label_ids)
540
        else:
541
            print("Wrong here1")
542

543
        if args.local_rank == -1:
544
            eval_sampler = RandomSampler(eval_data)
545
        else:
546
            eval_sampler = DistributedSampler(eval_data)
547
        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
548

549
        if mark:
550
            output_eval_file = os.path.join(args.output_dir, "eval_results_{}.txt".format(x.split("_")[-1]))
551
            output_file_pred = os.path.join(args.output_dir, "eval_pred_{}.txt".format(x.split("_")[-1]))
552
            output_file_glod = os.path.join(args.output_dir, "eval_gold_{}.txt".format(x.split("_")[-1]))
553
        else:
554
            output_eval_file = os.path.join(args.output_dir, "test_results_{}.txt".format(x.split("_")[-1]))
555
            output_file_pred = os.path.join(args.output_dir, "test_pred_{}.txt".format(x.split("_")[-1]))
556
            output_file_glod = os.path.join(args.output_dir, "test_gold_{}.txt".format(x.split("_")[-1]))
557

558
        fpred = open(output_file_pred, "w")
559
        fgold = open(output_file_glod, "w")
560

561
        model.eval()
562
        eval_loss, eval_accuracy = 0, 0
563
        nb_eval_steps, nb_eval_examples = 0, 0
564

565

566
        for step, batch in enumerate(tqdm(eval_dataloader, desc="Iteration")):
567
            #batch = tuple(t.to(device) if i != 3 else t for i, t in enumerate(batch))
568
            batch = tuple(t.to(device) for i, t in enumerate(batch))
569

570
            if args.task == 1:
571
                input_ids, attention_mask, label_ids = batch
572
            elif args.task == 2:
573
                input_ids, attention_mask, segment_ids, label_ids = batch
574
            else:
575
                print("Wrong here3")
576

577

578
            if args.task == 1:
579
                #loss, logits, hidden_states, attentions
580
                '''
581
                output = model(input_ids=input_ids, token_type_ids=None, attention_mask=attention_mask, labels=label_ids)
582
                logits = output.logits
583
                tmp_eval_loss = output.loss
584
                '''
585
                #
586
                tmp_eval_loss, logits = model(input_ids_org=input_ids, sentence_label=label_ids, attention_mask=attention_mask, func="task_class")
587
                #logits = output.logits
588
                #tmp_eval_loss = output.loss
589
            elif args.task == 2:
590
                #loss, logits, hidden_states, attentions
591
                '''
592
                output = model(input_ids=input_ids, token_type_ids=None, attention_mask=attention_mask, labels=label_ids)
593
                logits = output.logits
594
                tmp_eval_loss = output.loss
595
                '''
596
                #
597
                tmp_eval_loss, logits = model(input_ids_org=input_ids, sentence_label=label_ids, attention_mask=attention_mask, func="task_class")
598
                #exit()
599
                #logits = output.logits
600
                #tmp_eval_loss = output.loss
601
            else:
602
                print("Wrong!!")
603

604

605
            logits = logits.detach().cpu().numpy()
606
            label_ids = label_ids.to('cpu').numpy()
607
            tmp_eval_accuracy, pred = accuracy(logits, label_ids)
608
            for a, b in zip(pred, label_ids):
609
                fgold.write("{}\n".format(b))
610
                fpred.write("{}\n".format(a))
611

612
            eval_loss += tmp_eval_loss.mean().item()
613
            eval_accuracy += tmp_eval_accuracy
614

615
            nb_eval_examples += input_ids.size(0)
616
            nb_eval_steps += 1
617

618
        eval_loss = eval_loss / nb_eval_steps
619
        eval_accuracy = eval_accuracy / nb_eval_examples
620

621
        result = {'eval_loss': eval_loss,
622
                  'eval_accuracy': eval_accuracy
623
                  }
624

625
        with open(output_eval_file, "w") as writer:
626
            logger.info("***** Eval results *****")
627
            for key in sorted(result.keys()):
628
                logger.info("  %s = %s", key, str(result[key]))
629
                writer.write("%s = %s\n" % (key, str(result[key])))
630

631
        #if mark and step > int(math.ceil(len(eval_examples)/args.eval_batch_size)):
632
        if mark:
633
            model_performace_dev[x] = eval_accuracy
634
        else:
635
            model_performace_test[x] = eval_accuracy
636

637
    #################
638
    #################
639
    #####dev#########
640
    if args.choose_eval_test_both != 1:
641
        model_name_best=0
642
        score_best=0
643
        for model_name, score in model_performace_dev.items():
644
            if score >= score_best:
645
                score_best = score
646
                model_name_best = model_name
647

648

649
        model = RobertaForMaskedLMDomainTask.from_pretrained(args.pretrain_model, output_hidden_states=False, output_attentions=False, return_dict=True, num_labels=args.num_labels_task)
650
        model_name_best = os.path.join(args.output_dir, model_name_best)
651
        model.load_state_dict(torch.load(model_name_best), strict=False)
652
        # Save a trained model
653
        logger.info("** ** * Saving fine - tuned model ** ** * ")
654
        model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
655
        output_model_file = os.path.join(args.output_dir, "pytorch_model.bin_dev_best")
656
        torch.save(model_to_save.state_dict(), output_model_file)
657

658

659
    if args.choose_eval_test_both != 0:
660
        model_name_best=0
661
        score_best=0
662
        for model_name, score in model_performace_test.items():
663
            if score >= score_best:
664
                score_best = score
665
                model_name_best = model_name
666

667

668
        model = RobertaForMaskedLMDomainTask.from_pretrained(args.pretrain_model, output_hidden_states=False, output_attentions=False, return_dict=True, num_labels=args.num_labels_task)
669
        model_name_best = os.path.join(args.output_dir, model_name_best)
670
        model.load_state_dict(torch.load(model_name_best), strict=False)
671
        # Save a trained model
672
        logger.info("** ** * Saving fine - tuned model ** ** * ")
673
        model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
674
        output_model_file = os.path.join(args.output_dir, "pytorch_model.bin_test_best")
675
        torch.save(model_to_save.state_dict(), output_model_file)
676

677
if __name__ == "__main__":
678
    main()
679

680

681