| 40 | |
| 41 | |
| 42 | class Bert(nn.Module): |
| 43 | def __init__(self, model_dir, temp_dir, load_pretrained_bert, bert_config): |
| 44 | super(Bert, self).__init__() |
| 45 | if load_pretrained_bert: |
| 46 | # self.models = BertModel.from_pretrained('bert-base-uncased', cache_dir=temp_dir) |
| 47 | self.model = BertModel.from_pretrained(model_dir, cache_dir=temp_dir) |
| 48 | else: |
| 49 | self.model = BertModel(bert_config) |
| 50 | |
| 51 | def forward(self, x, segs, mask): |
| 52 | encoded_layers, _ = self.model(x, segs, attention_mask=mask) |
| 53 | top_vec = encoded_layers[-1] |
| 54 | return top_vec |
| 55 | |
| 56 | |
| 57 | class Summarizer(nn.Module): |
no outgoing calls
no test coverage detected