File size: 2,135 Bytes
4d5aff4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import torch
from transformers import GPT2ForSequenceClassification
class ClassificationHead(torch.nn.Module):
"""Classification Head for transformer encoders"""
def __init__(self, class_size, embed_size, is_deep=False, use_xlnet=False, is_deeper=False):
super(ClassificationHead, self).__init__()
self.class_size = class_size
self.embed_size = embed_size
self.is_deep = is_deep
self.is_deeper = is_deeper
self.use_xlnet = use_xlnet
if is_deep:
self.mlp1 = torch.nn.Linear(embed_size, 128)
self.mlp2 = torch.nn.Linear(128, 64)
self.mlp3 = torch.nn.Linear(64, class_size)
elif is_deeper:
self.mlp1 = torch.nn.Linear(embed_size, 512)
self.mlp2 = torch.nn.Linear(512, 256)
self.mlp3 = torch.nn.Linear(256, 128)
self.mlp4 = torch.nn.Linear(128, 64)
self.mlp5 = torch.nn.Linear(64, class_size)
elif use_xlnet:
self.gpt = GPT2ForSequenceClassification.from_pretrained("microsoft/DialogRPT-updown")
self.mlp = torch.nn.Linear(8, class_size, bias=True)
else:
self.mlp = torch.nn.Linear(embed_size, class_size)
def forward(self, hidden_state, inputs_embeds=None):
if self.is_deep:
hidden_state = torch.nn.functional.relu(self.mlp1(hidden_state))
hidden_state = torch.nn.functional.relu(self.mlp2(hidden_state))
logits = self.mlp3(hidden_state)
elif self.is_deeper:
hidden_state = torch.nn.functional.relu(self.mlp1(hidden_state))
hidden_state = torch.nn.functional.relu(self.mlp2(hidden_state))
hidden_state = torch.nn.functional.relu(self.mlp3(hidden_state))
hidden_state = torch.nn.functional.relu(self.mlp4(hidden_state))
logits = self.mlp5(hidden_state)
elif self.use_xlnet:
hidden_state, _ = self.gpt(input_ids=hidden_state, inputs_embeds=inputs_embeds)
logits = self.mlp(hidden_state)
else:
logits = self.mlp(hidden_state)
return logits
|