debajyotimaz commited on
Commit
aa14a33
·
verified ·
1 Parent(s): 89b78d4

Upload models.py

Browse files
Files changed (1) hide show
  1. models.py +50 -0
models.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import BertTokenizer
3
+ from torch import nn
4
+ from transformers import BertModel
5
+
6
+
7
+ #CustomMBERTModel= torch.load("/data2/Akash_for_interface/model_mbert_1416.pt")
8
+ tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased')
9
+ sentences= "you are good person."
10
+ max_len=len(sentences)
11
+ encoding = tokenizer.encode_plus(sentences,add_special_tokens=True,max_length=max_len,padding='max_length',truncation=True,return_tensors='pt')
12
+ input_ids=encoding['input_ids'].flatten(),
13
+ attention_mask= encoding['attention_mask'].flatten()
14
+ #print(input_ids[0])
15
+
16
+ labels=["Non Hateful","Hateful"]
17
+
18
+
19
+
20
+ device = 'cpu'
21
+ class CustomMBERTModel(nn.Module):
22
+ def __init__(self, num_labels):
23
+ super(CustomMBERTModel, self).__init__()
24
+ self.bert = BertModel.from_pretrained('bert-base-multilingual-cased')
25
+
26
+ # Freeze all layers except the top 2
27
+ for param in self.bert.parameters():
28
+ param.requires_grad = False
29
+
30
+ # Unfreeze the parameters of the top 2 layers
31
+ for param in self.bert.encoder.layer[-2:].parameters():
32
+ param.requires_grad = True
33
+
34
+ # Adding Linear layer
35
+ self.linear = nn.Linear(self.bert.config.hidden_size, num_labels)
36
+
37
+ def forward(self, input_ids, attention_mask):
38
+ outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
39
+ hidden_states = outputs.last_hidden_state
40
+
41
+ pooled_output= hidden_states[:,0,:]
42
+ pooled_output=torch.squeeze(pooled_output,dim=1)
43
+
44
+ #print('p-shape:', pooled_output.shape)
45
+
46
+ logits = self.linear(pooled_output)
47
+
48
+ #print('l-shape:', logits.shape)
49
+ return logits
50
+