haritzpuerto commited on
Commit
dd38fc4
1 Parent(s): 4b2efc2

Create model.py

Browse files
Files changed (1) hide show
  1. model.py +301 -0
model.py ADDED
@@ -0,0 +1,301 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from torch import nn
2
+ import torch
3
+ import numpy as np
4
+
5
+ from transformers import BertPreTrainedModel
6
+ from transformers.modeling_outputs import TokenClassifierOutput, BaseModelOutputWithPoolingAndCrossAttentions
7
+ from transformers.models.bert.modeling_bert import BertPooler, BertEncoder
8
+
9
+ class MetaQA_Model(BertPreTrainedModel):
10
+ def __init__(self, config):
11
+ super().__init__(config)
12
+ self.bert = MetaQABertModel(config)
13
+ self.num_agents = config.num_agents
14
+
15
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
16
+ self.list_MoSeN = nn.ModuleList([nn.Linear(config.hidden_size, 1) for i in range(self.num_agents)])
17
+ self.input_size_ans_sel = 1 + config.hidden_size
18
+ interm_size = int(config.hidden_size/2)
19
+ self.ans_sel = nn.Sequential(nn.Linear(self.input_size_ans_sel, interm_size),
20
+ nn.ReLU(),
21
+ nn.Dropout(config.hidden_dropout_prob),
22
+ nn.Linear(interm_size, 2))
23
+
24
+ self.init_weights()
25
+
26
+ def forward(
27
+ self,
28
+ input_ids=None,
29
+ attention_mask=None,
30
+ token_type_ids=None,
31
+ position_ids=None,
32
+ head_mask=None,
33
+ inputs_embeds=None,
34
+ labels=None,
35
+ output_attentions=None,
36
+ output_hidden_states=None,
37
+ return_dict=None,
38
+ ans_sc=None,
39
+ agent_sc=None,
40
+ ):
41
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
42
+
43
+ outputs = self.bert(
44
+ input_ids,
45
+ attention_mask=attention_mask,
46
+ token_type_ids=token_type_ids,
47
+ position_ids=position_ids,
48
+ head_mask=head_mask,
49
+ inputs_embeds=inputs_embeds,
50
+ output_attentions=output_attentions,
51
+ output_hidden_states=output_hidden_states,
52
+ return_dict=return_dict,
53
+ ans_sc=ans_sc,
54
+ agent_sc=agent_sc,
55
+ )
56
+ # domain classification
57
+ pooled_output = outputs[1]
58
+
59
+ pooled_output = self.dropout(pooled_output)
60
+ list_domains_logits = []
61
+ for MoSeN in self.list_MoSeN:
62
+ domain_logits = MoSeN(pooled_output)
63
+ list_domains_logits.append(domain_logits)
64
+ domain_logits = torch.stack(list_domains_logits)
65
+ # shape = (num_agents, batch_size, 1)
66
+ # we have to transpose the shape to (batch_size, num_agents, 1)
67
+ domain_logits = domain_logits.transpose(0,1)
68
+
69
+ # ans classifier
70
+ sequence_output = outputs[0] # (batch_size, seq_len, hidden_size)
71
+ # select the [RANK] token embeddings
72
+ idx_rank = (input_ids == 1).nonzero() # (batch_size x num_agents, 2)
73
+ idx_rank = idx_rank[:,1].view(-1, self.num_agents)
74
+ list_emb = []
75
+ for i in range(idx_rank.shape[0]):
76
+ rank_emb = sequence_output[i][idx_rank[i], :]
77
+ # rank shape = (1, hidden_size)
78
+ list_emb.append(rank_emb)
79
+
80
+ rank_emb = torch.stack(list_emb)
81
+
82
+ rank_emb = self.dropout(rank_emb)
83
+ rank_emb = torch.cat((rank_emb, domain_logits), dim=2)
84
+ # rank emb shape = (batch_size, num_agents, hidden_size+1)
85
+ logits = self.ans_sel(rank_emb) # (batch_size, num_agents, 2)
86
+
87
+ if not return_dict:
88
+ output = (logits,) + outputs[2:]
89
+ return output
90
+
91
+ return TokenClassifierOutput(
92
+ loss=None,
93
+ logits=logits,
94
+ hidden_states=outputs.hidden_states,
95
+ attentions=outputs.attentions,
96
+ )
97
+
98
+
99
+ class MetaQABertModel(BertPreTrainedModel):
100
+ def __init__(self, config, add_pooling_layer=True):
101
+ super().__init__(config)
102
+ self.config = config
103
+
104
+ self.embeddings = MetaQABertEmbeddings(config) # NEW
105
+ self.encoder = BertEncoder(config)
106
+ self.pooler = BertPooler(config) if add_pooling_layer else None
107
+
108
+ self.init_weights()
109
+
110
+ def get_input_embeddings(self):
111
+ return self.embeddings.word_embeddings
112
+
113
+ def set_input_embeddings(self, value):
114
+ self.embeddings.word_embeddings = value
115
+
116
+ def _prune_heads(self, heads_to_prune):
117
+ for layer, heads in heads_to_prune.items():
118
+ self.encoder.layer[layer].attention.prune_heads(heads)
119
+
120
+ def forward(
121
+ self,
122
+ input_ids=None,
123
+ attention_mask=None,
124
+ token_type_ids=None,
125
+ position_ids=None,
126
+ head_mask=None,
127
+ inputs_embeds=None,
128
+ encoder_hidden_states=None,
129
+ encoder_attention_mask=None,
130
+ past_key_values=None,
131
+ use_cache=None,
132
+ output_attentions=None,
133
+ output_hidden_states=None,
134
+ return_dict=None,
135
+ ans_sc=None,
136
+ agent_sc=None,
137
+ ):
138
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
139
+ output_hidden_states = (
140
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
141
+ )
142
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
143
+
144
+ if self.config.is_decoder:
145
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
146
+ else:
147
+ use_cache = False
148
+
149
+ if input_ids is not None and inputs_embeds is not None:
150
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
151
+ elif input_ids is not None:
152
+ input_shape = input_ids.size()
153
+ batch_size, seq_length = input_shape
154
+ elif inputs_embeds is not None:
155
+ input_shape = inputs_embeds.size()[:-1]
156
+ batch_size, seq_length = input_shape
157
+ else:
158
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
159
+
160
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
161
+
162
+ # past_key_values_length
163
+ past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
164
+
165
+ if attention_mask is None:
166
+ attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
167
+
168
+ if token_type_ids is None:
169
+ if hasattr(self.embeddings, "token_type_ids"):
170
+ buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
171
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
172
+ token_type_ids = buffered_token_type_ids_expanded
173
+ else:
174
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
175
+
176
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
177
+ # ourselves in which case we just need to make it broadcastable to all heads.
178
+ extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)
179
+
180
+ # If a 2D or 3D attention mask is provided for the cross-attention
181
+ # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
182
+ if self.config.is_decoder and encoder_hidden_states is not None:
183
+ encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
184
+ encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
185
+ if encoder_attention_mask is None:
186
+ encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
187
+ encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
188
+ else:
189
+ encoder_extended_attention_mask = None
190
+
191
+ # Prepare head mask if needed
192
+ # 1.0 in head_mask indicate we keep the head
193
+ # attention_probs has shape bsz x n_heads x N x N
194
+ # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
195
+ # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
196
+ head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
197
+
198
+ embedding_output = self.embeddings(
199
+ input_ids=input_ids,
200
+ position_ids=position_ids,
201
+ token_type_ids=token_type_ids,
202
+ inputs_embeds=inputs_embeds,
203
+ past_key_values_length=past_key_values_length,
204
+ ans_sc=ans_sc,
205
+ agent_sc=agent_sc,
206
+ )
207
+ encoder_outputs = self.encoder(
208
+ embedding_output,
209
+ attention_mask=extended_attention_mask,
210
+ head_mask=head_mask,
211
+ encoder_hidden_states=encoder_hidden_states,
212
+ encoder_attention_mask=encoder_extended_attention_mask,
213
+ past_key_values=past_key_values,
214
+ use_cache=use_cache,
215
+ output_attentions=output_attentions,
216
+ output_hidden_states=output_hidden_states,
217
+ return_dict=return_dict,
218
+ )
219
+ sequence_output = encoder_outputs[0]
220
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
221
+
222
+ if not return_dict:
223
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
224
+
225
+ return BaseModelOutputWithPoolingAndCrossAttentions(
226
+ last_hidden_state=sequence_output,
227
+ pooler_output=pooled_output,
228
+ past_key_values=encoder_outputs.past_key_values,
229
+ hidden_states=encoder_outputs.hidden_states,
230
+ attentions=encoder_outputs.attentions,
231
+ cross_attentions=encoder_outputs.cross_attentions,
232
+ )
233
+
234
+ class MetaQABertEmbeddings(nn.Module):
235
+ """Construct the embeddings from
236
+ word, position, token_type embeddings, and scores from the QA agents."""
237
+
238
+ def __init__(self, config):
239
+ super().__init__()
240
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
241
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
242
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
243
+ self.ans_sc_proj = nn.Linear(1, config.hidden_size)
244
+ self.agent_sc_proj = nn.Linear(1, config.hidden_size)
245
+ # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
246
+ # any TensorFlow checkpoint file
247
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
248
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
249
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
250
+ self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
251
+ self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
252
+ self.register_buffer(
253
+ "token_type_ids",
254
+ torch.zeros(self.position_ids.size(), dtype=torch.long, device=self.position_ids.device),
255
+ persistent=False,
256
+ )
257
+
258
+
259
+ def forward(
260
+ self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0,
261
+ ans_sc=None, agent_sc=None):
262
+ if input_ids is not None:
263
+ input_shape = input_ids.size()
264
+ else:
265
+ input_shape = inputs_embeds.size()[:-1]
266
+
267
+ seq_length = input_shape[1]
268
+
269
+ if position_ids is None:
270
+ position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
271
+
272
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
273
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
274
+ # issue #5664
275
+ if token_type_ids is None:
276
+ if hasattr(self, "token_type_ids"):
277
+ buffered_token_type_ids = self.token_type_ids[:, :seq_length]
278
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
279
+ token_type_ids = buffered_token_type_ids_expanded
280
+ else:
281
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
282
+
283
+ if inputs_embeds is None:
284
+ inputs_embeds = self.word_embeddings(input_ids)
285
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
286
+
287
+ embeddings = inputs_embeds + token_type_embeddings
288
+ if self.position_embedding_type == "absolute":
289
+ position_embeddings = self.position_embeddings(position_ids)
290
+ embeddings += position_embeddings
291
+
292
+ if ans_sc is not None:
293
+ ans_sc_emb = self.ans_sc_proj(ans_sc.unsqueeze(2))
294
+ embeddings += ans_sc_emb
295
+ if agent_sc is not None:
296
+ agent_sc_emb = self.agent_sc_proj(agent_sc.unsqueeze(2))
297
+ embeddings += agent_sc_emb
298
+
299
+ embeddings = self.LayerNorm(embeddings)
300
+ embeddings = self.dropout(embeddings)
301
+ return embeddings