class Discriminator(torch.nn.Module):
"""Transformer encoder followed by a Classification Head"""
def __init__(
self,
class_size,
pretrained_model="gpt2-medium",
cached_mode=False,
device='cpu'
):
super(Discriminator, self).__init__()
self.tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model)
self.encoder = GPT2LMHeadModel.from_pretrained(pretrained_model)
self.embed_size = self.encoder.transformer.config.hidden_size
self.classifier_head = ClassificationHead(
class_size=class_size,
embed_size=self.embed_size
)
self.cached_mode = cached_mode
self.device = device
def get_classifier(self):
return self.classifier_head
def train_custom(self):
for param in self.encoder.parameters():
param.requires_grad = False
self.classifier_head.train()
def avg_representation(self, x):
mask = x.ne(0).unsqueeze(2).repeat(
1, 1, self.embed_size
).float().to(self.device).detach()
hidden, _ = self.encoder.transformer(x)
masked_hidden = hidden * mask
avg_hidden = torch.sum(masked_hidden, dim=1) / (
torch.sum(mask, dim=1).detach() + EPSILON
)
return avg_hidden
def forward(self, x):
if self.cached_mode:
avg_hidden = x.to(self.device)
else:
avg_hidden = self.avg_representation(x.to(self.device))
logits = self.classifier_head(avg_hidden)
probs = F.log_softmax(logits, dim=-1)
return probs
class Discriminator(torch.nn.Module):
"""Transformer encoder followed by a Classification Head"""
def __init__(self,
class_size=None,
pretrained_model="gpt2-medium",
classifier_head=None,
cached_mode=False,
device=DEVICE):
super(Discriminator, self).__init__()
if pretrained_model.startswith("gpt2"):
self.tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model)
self.encoder = GPT2LMHeadModel.from_pretrained(pretrained_model)
self.embed_size = self.encoder.transformer.config.hidden_size
elif pretrained_model.startswith("bert"):
self.tokenizer = BertTokenizer.from_pretrained(pretrained_model)
self.encoder = BertModel.from_pretrained(pretrained_model)
self.embed_size = self.encoder.config.hidden_size
else:
raise ValueError(
"{} model not yet supported".format(pretrained_model))
if classifier_head:
self.classifier_head = classifier_head
else:
if not class_size:
raise ValueError("must specify class_size")
self.classifier_head = ClassificationHead(
class_size=class_size, embed_size=self.embed_size)
self.cached_mode = cached_mode
self.device = device
def get_classifier(self):
return self.classifier_head
def train_custom(self):
for param in self.encoder.parameters():
param.requires_grad = False
self.classifier_head.train()
def avg_representation(self, x):
mask = x.ne(0).unsqueeze(2).repeat(1, 1, self.embed_size).float().to(
self.device).detach()
if hasattr(self.encoder, 'transformer'):
# for gpt2
hidden, _ = self.encoder.transformer(x)
else:
# for bert
hidden, _ = self.encoder(x)
masked_hidden = hidden * mask
avg_hidden = torch.sum(
masked_hidden, dim=1) / (torch.sum(mask, dim=1).detach() + EPSILON)
return avg_hidden
def forward(self, x):
if self.cached_mode:
avg_hidden = x.to(self.device)
else:
avg_hidden = self.avg_representation(x.to(self.device))
logits = self.classifier_head(avg_hidden)
probs = F.log_softmax(logits, dim=-1)
return probs
def predict(self, input_sentence):
input_t = self.tokenizer.encode(input_sentence)
input_t = torch.tensor([input_t], dtype=torch.long, device=self.device)
if self.cached_mode:
input_t = self.avg_representation(input_t)
log_probs = self(input_t).data.cpu().numpy().flatten().tolist()
prob = [math.exp(log_prob) for log_prob in log_probs]
return prob