(self, args, device, load_pretrained_bert=False, bert_config=None)
| 56 | |
| 57 | class Summarizer(nn.Module): |
| 58 | def __init__(self, args, device, load_pretrained_bert=False, bert_config=None): |
| 59 | super(Summarizer, self).__init__() |
| 60 | self.args = args |
| 61 | self.device = device |
| 62 | self.bert = Bert(args.bert_pretrained_model_path, args.temp_dir, load_pretrained_bert, bert_config) |
| 63 | if args.encoder == 'classifier': |
| 64 | self.encoder = Classifier(self.bert.model.config.hidden_size) |
| 65 | # TODO: 加入了返回 sentence vector |
| 66 | elif args.encoder == 'transformer': |
| 67 | self.encoder = TransformerInterEncoder(self.bert.model.config.hidden_size, args.ff_size, args.heads, |
| 68 | args.dropout, args.inter_layers) |
| 69 | elif args.encoder == 'rnn': |
| 70 | self.encoder = RNNEncoder(bidirectional=True, num_layers=1, |
| 71 | input_size=self.bert.model.config.hidden_size, hidden_size=args.rnn_size, |
| 72 | dropout=args.dropout) |
| 73 | elif args.encoder == 'baseline': |
| 74 | bert_config = BertConfig(self.bert.model.config.vocab_size, hidden_size=args.hidden_size, |
| 75 | num_hidden_layers=6, num_attention_heads=8, intermediate_size=args.ff_size) |
| 76 | self.bert.model = BertModel(bert_config) |
| 77 | self.encoder = Classifier(self.bert.model.config.hidden_size) |
| 78 | |
| 79 | if args.param_init != 0.0: |
| 80 | for p in self.encoder.parameters(): |
| 81 | p.data.uniform_(-args.param_init, args.param_init) |
| 82 | if args.param_init_glorot: |
| 83 | for p in self.encoder.parameters(): |
| 84 | if p.dim() > 1: |
| 85 | xavier_uniform_(p) |
| 86 | |
| 87 | self.to(device) |
| 88 | |
| 89 | def load_cp(self, pt): |
| 90 | self.load_state_dict(pt['models'], strict=True) |
nothing calls this directly
no test coverage detected