| 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)) |
| |
