Views
No views yet
1from transformers import BertModel
2
3class BertClassifier(nn.Module, PyTorchModelHubMixin):
4 def __init__(self, dataset: str, num_classes, dropout=0.5):
5
6 super(BertClassifier, self).__init__()
7
8 self.model_name = "bert-base-uncased"
9 print(f"Loading BERT model {self.model_name} for {dataset} dataset...")
10
11 self.bert = BertModel.from_pretrained(self.model_name)
12 self.dropout = nn.Dropout(dropout)
13 self.linear = nn.Linear(768, num_classes) # in features, out features = number of classes
14 self.relu = nn.ReLU()
15
16 def forward(self, input_ids, attention_mask):
17
18 _, pooled_output = self.bert(input_ids=input_ids, attention_mask=attention_mask, return_dict=False)
19 dropout_output = self.dropout(pooled_output)
20 linear_output = self.linear(dropout_output)
21 final_layer = self.relu(linear_output)
22
23 return final_layermodel = BertClassifier.from_pretrained("CDL-RecSys/BERT-uncased-fic-category-classifier")BertTokenizer:1from transformers import BertModel, BertTokenizer
2
3tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
4input = "this shawl collar jersey blazer bloom with black and white floral patterning inspired by the work of rising spanish photographer coco capit n"
5
6texts = self.tokenizer(batch, padding='max_length', max_length = 512, truncation=True,return_tensors="pt")
7input_ids = texts["input_ids"]
8attention_mask = texts["attention_mask"]
9
10output = model(input_ids, attention_mask)
11class = output.argmax(dim=1) # should be 51 (blazer)