forked from YJiangcm/SST-2-sentiment-analysis
-
Notifications
You must be signed in to change notification settings - Fork 0
/
models.py
87 lines (70 loc) · 3.92 KB
/
models.py
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
# -*- coding: utf-8 -*-
"""
Created on Sun Oct 25 00:19:30 2020
@author: Jiang Yuxin
"""
import torch
from torch import nn
from transformers import (
BertForSequenceClassification,
AlbertForSequenceClassification,
XLNetForSequenceClassification,
RobertaForSequenceClassification,
AutoTokenizer
)
class AlbertModel(nn.Module):
def __init__(self, requires_grad = True):
super(AlbertModel, self).__init__()
self.albert = AlbertForSequenceClassification.from_pretrained('albert-xxlarge-v2', num_labels = 2)
self.tokenizer = AutoTokenizer.from_pretrained('albert-xxlarge-v2', do_lower_case=True)
self.requires_grad = requires_grad
self.device = torch.device("cuda")
for param in self.albert.parameters():
param.requires_grad = True # Each parameter requires gradient
def forward(self, batch_seqs, batch_seq_masks, batch_seq_segments, labels):
loss, logits = self.albert(input_ids = batch_seqs, attention_mask = batch_seq_masks,
token_type_ids=batch_seq_segments, labels = labels)[:2]
probabilities = nn.functional.softmax(logits, dim=-1)
return loss, logits, probabilities
class BertModel(nn.Module):
def __init__(self, requires_grad = True):
super(BertModel, self).__init__()
self.bert = BertForSequenceClassification.from_pretrained('textattack/bert-base-uncased-SST-2',num_labels = 2)
self.tokenizer = AutoTokenizer.from_pretrained('textattack/bert-base-uncased-SST-2', do_lower_case=True)
self.requires_grad = requires_grad
self.device = torch.device("cuda")
for param in self.bert.parameters():
param.requires_grad = requires_grad # Each parameter requires gradient
def forward(self, batch_seqs, batch_seq_masks, batch_seq_segments, labels):
loss, logits = self.bert(input_ids = batch_seqs, attention_mask = batch_seq_masks,
token_type_ids=batch_seq_segments, labels = labels)[:2]
probabilities = nn.functional.softmax(logits, dim=-1)
return loss, logits, probabilities
class RobertModel(nn.Module):
def __init__(self, requires_grad = True):
super(RobertModel, self).__init__()
self.bert = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels = 2)
self.tokenizer = AutoTokenizer.from_pretrained('roberta-base', do_lower_case=True)
self.requires_grad = requires_grad
self.device = torch.device("cuda")
for param in self.bert.parameters():
param.requires_grad = requires_grad # Each parameter requires gradient
def forward(self, batch_seqs, batch_seq_masks, batch_seq_segments, labels):
loss, logits = self.bert(input_ids = batch_seqs, attention_mask = batch_seq_masks,
token_type_ids=batch_seq_segments, labels = labels)[:2]
probabilities = nn.functional.softmax(logits, dim=-1)
return loss, logits, probabilities
class XlnetModel(nn.Module):
def __init__(self, requires_grad = True):
super(XlnetModel, self).__init__()
self.xlnet = XLNetForSequenceClassification.from_pretrained('xlnet-large-cased', num_labels = 2)
self.tokenizer = AutoTokenizer.from_pretrained('xlnet-large-cased', do_lower_case=True)
self.requires_grad = requires_grad
self.device = torch.device("cuda")
for param in self.xlnet.parameters():
param.requires_grad = requires_grad # Each parameter requires gradient
def forward(self, batch_seqs, batch_seq_masks, batch_seq_segments, labels):
loss, logits = self.xlnet(input_ids = batch_seqs, attention_mask = batch_seq_masks,
token_type_ids=batch_seq_segments, labels = labels)[:2]
probabilities = nn.functional.softmax(logits, dim=-1)
return loss, logits, probabilities