|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
from torch.nn.utils.rnn import pad_sequence |
|
from transformers import AutoModel |
|
from pathlib import Path |
|
|
|
|
|
class LinearTokenSelector(nn.Module): |
|
def __init__(self, encoder, embedding_size=768): |
|
super(LinearTokenSelector, self).__init__() |
|
self.encoder = encoder |
|
self.classifier = nn.Linear(embedding_size, 2, bias=False) |
|
|
|
def forward(self, x): |
|
output = self.encoder(x, output_hidden_states=True) |
|
x = output["hidden_states"][-1] |
|
x = self.classifier(x) |
|
x = F.log_softmax(x, dim=2) |
|
return x |
|
|
|
def save(self, classifier_path, encoder_path): |
|
state = self.state_dict() |
|
state = dict((k, v) for k, v in state.items() if k.startswith("classifier")) |
|
torch.save(state, classifier_path) |
|
self.encoder.save_pretrained(encoder_path) |
|
|
|
def predict(self, texts, tokenizer, device): |
|
input_ids = tokenizer(texts)["input_ids"] |
|
input_ids = pad_sequence( |
|
[torch.tensor(ids) for ids in input_ids], batch_first=True |
|
).to(device) |
|
logits = self.forward(input_ids) |
|
argmax_labels = torch.argmax(logits, dim=2) |
|
return labels_to_summary(input_ids, argmax_labels, tokenizer) |
|
|
|
|
|
def load_model(model_dir, device="cuda", prefix="best"): |
|
if isinstance(model_dir, str): |
|
model_dir = Path(model_dir) |
|
for p in (model_dir / "checkpoints").iterdir(): |
|
if p.name.startswith(f"{prefix}"): |
|
checkpoint_dir = p |
|
return load_checkpoint(checkpoint_dir, device=device) |
|
|
|
|
|
def load_checkpoint(checkpoint_dir, device="cuda"): |
|
if isinstance(checkpoint_dir, str): |
|
checkpoint_dir = Path(checkpoint_dir) |
|
|
|
encoder_path = checkpoint_dir / "encoder.bin" |
|
classifier_path = checkpoint_dir / "classifier.bin" |
|
|
|
encoder = AutoModel.from_pretrained(encoder_path).to(device) |
|
embedding_size = encoder.state_dict()["embeddings.word_embeddings.weight"].shape[1] |
|
|
|
classifier = LinearTokenSelector(None, embedding_size).to(device) |
|
classifier_state = torch.load(classifier_path, map_location=device) |
|
classifier_state = dict( |
|
(k, v) for k, v in classifier_state.items() |
|
if k.startswith("classifier") |
|
) |
|
classifier.load_state_dict(classifier_state) |
|
classifier.encoder = encoder |
|
return classifier.to(device) |
|
|
|
|
|
def labels_to_summary(input_batch, label_batch, tokenizer): |
|
summaries = [] |
|
for input_ids, labels in zip(input_batch, label_batch): |
|
selected = [int(input_ids[i]) for i in range(len(input_ids)) |
|
if labels[i] == 1] |
|
summary = tokenizer.decode(selected, skip_special_tokens=True) |
|
summaries.append(summary) |
|
return summaries |
|
|
