Spaces:
Sleeping
Sleeping
| from transformers import AutoTokenizer, AutoConfig, AutoModel | |
| import torch | |
| class CustomModel(torch.nn.Module): | |
| """ | |
| This takes a transformer backbone and puts a slightly-modified classification head on top. | |
| """ | |
| def __init__(self, model_name, num_extra_dims, num_labels=2): | |
| # num_extra_dims corresponds to the number of extra dimensions of numerical/categorical data | |
| super().__init__() | |
| self.config = AutoConfig.from_pretrained(model_name, num_labels=num_labels) | |
| self.transformer = AutoModel.from_pretrained(model_name, config=self.config) | |
| num_hidden_size = self.transformer.config.hidden_size # May be different depending on which model you use. Common sizes are 768 and 1024. Look in the config.json file | |
| self.linear_layer_1 = torch.nn.Linear(num_hidden_size+num_extra_dims, 32) | |
| # Output size is 1 since this is a binary classification problem | |
| self.linear_layer_2 = torch.nn.Linear(32, 16) | |
| self.linear_layer_output = torch.nn.Linear(16, 1) | |
| self.relu = torch.nn.LeakyReLU(0.6) | |
| self.dropout_1 = torch.nn.Dropout(0.5) | |
| def forward(self, input_ids, extra_features, attention_mask=None, token_type_ids=None, labels=None): | |
| """ | |
| extra_features should be of shape [batch_size, dim] | |
| where dim is the number of additional numerical/categorical dimensions | |
| """ | |
| hidden_states = self.transformer(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids) # [batch size, sequence length, hidden size] | |
| cls_embeds = hidden_states.last_hidden_state[:, 0, :] # [batch size, hidden size] | |
| concat = torch.cat((cls_embeds, extra_features), dim=-1) # [batch size, hidden size+num extra dims] | |
| output_1 = self.relu(self.linear_layer_1(concat)) # [batch size, num labels] | |
| output_2 = self.relu(self.linear_layer_2(output_1)) | |
| final_output = self.dropout_1(self.linear_layer_output(output_2)) | |
| return final_output |