|
|
import torch |
|
|
import torch.nn as nn |
|
|
from tqdm import tqdm |
|
|
from transformers import DistilBertTokenizerFast, DistilBertModel |
|
|
import numpy as np |
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
|
tokenizer = DistilBertTokenizerFast.from_pretrained("distilbert-base-uncased") |
|
|
|
|
|
class DistilBERTSent(nn.Module): |
|
|
""" |
|
|
DistilBERT but with a layer attached to perform binary classification. |
|
|
""" |
|
|
def __init__(self, freeze_bert=False): |
|
|
super(DistilBERTSent, self).__init__() |
|
|
self.distil_bert = DistilBertModel.from_pretrained('distilbert-base-uncased') |
|
|
self.linear = nn.Linear(2304, 1) |
|
|
self.sigmoid = nn.Sigmoid() |
|
|
|
|
|
if freeze_bert: |
|
|
for param in self.distil_bert.parameters(): |
|
|
param.requires_grad = False |
|
|
|
|
|
def forward(self, ids, mask): |
|
|
outputs = self.distil_bert(input_ids = ids, attention_mask=mask, output_hidden_states=True) |
|
|
x = torch.concat(outputs.hidden_states[:-4], dim=2).mean(1) |
|
|
x = self.linear(x) |
|
|
x = self.sigmoid(x) |
|
|
return x.flatten() |
|
|
|
|
|
def initialize(path="models/model.pt"): |
|
|
model = DistilBERTSent() |
|
|
model.load_state_dict(torch.load(path, map_location=device)) |
|
|
model.to(device) |
|
|
model.eval() |
|
|
return model |
|
|
|
|
|
def chunks(lst, n): |
|
|
for i in tqdm(range(0, len(lst), n)): |
|
|
yield lst[i:i+n] |
|
|
|
|
|
@torch.no_grad() |
|
|
def inference(model, text, batch_size=32): |
|
|
""" |
|
|
pass in model, list of text, and batch_size |
|
|
""" |
|
|
to_return = [] |
|
|
for batch in chunks(text, batch_size): |
|
|
encoded = tokenizer( |
|
|
text = batch, |
|
|
add_special_tokens=True, |
|
|
padding='max_length', |
|
|
return_attention_mask=True, |
|
|
truncation=True |
|
|
) |
|
|
input_ids = torch.tensor(encoded.get('input_ids')).to(device) |
|
|
attention_masks = torch.tensor(encoded.get('attention_mask')).to(device) |
|
|
to_return.append(model(input_ids, attention_masks)) |
|
|
|
|
|
return torch.concat(to_return).cpu().numpy() |
|
|
|
|
|
if __name__ == "__main__": |
|
|
model = initialize() |
|
|
text = ["I love it so much!", "Broke on the first day"] |
|
|
print(inference(model, text, 2)) |
|
|
|
