|
from typing import List |
|
|
|
import pandas as pd |
|
import torch |
|
from tqdm import trange |
|
from transformers import AutoModel, AutoTokenizer |
|
|
|
from src.feature_extractors.base_extractor import BaseExtractor |
|
|
|
|
|
class BertPretrainFeatureExtractor(BaseExtractor): |
|
"""Extract [CLS] embedding feature from any untrained bert-like models""" |
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
def __init__(self, model_name: str, max_length: int = 512, batch_size=64): |
|
self.model_name = model_name |
|
self.max_length = max_length |
|
self.batch_size = batch_size |
|
|
|
self.model = AutoModel.from_pretrained(self.model_name) |
|
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) |
|
|
|
@torch.no_grad() |
|
def generate_features(self, data: pd.Series) -> pd.DataFrame: |
|
""" |
|
Generates features in batch-mode, obtained from untrained bert model. |
|
|
|
:param data: Series with full_text column |
|
:return: Dataframe, that have index - id's from data, and columns - bert features |
|
""" |
|
torch.cuda.empty_cache() |
|
|
|
texts = data.tolist() |
|
self.model = self.model.to(self.device) |
|
|
|
classification_outputs = [] |
|
for ii in trange( |
|
0, len(data), self.batch_size, |
|
total=len(data) // self.batch_size + 1, |
|
desc="Generating bert features..." |
|
): |
|
text_batch = texts[ii: ii + self.batch_size] |
|
batch_encoded = self.tokenizer( |
|
text_batch, |
|
padding=True, |
|
truncation=True, |
|
max_length=self.max_length, |
|
return_tensors='pt' |
|
).to(self.device) |
|
|
|
output = self.model(**batch_encoded) |
|
cls_output = output['last_hidden_state'][:, 0].cpu() |
|
classification_outputs.append(cls_output) |
|
|
|
self.model = self.model.to("cpu") |
|
classification_outputs_tensor = torch.cat(classification_outputs, dim=0) |
|
torch.cuda.empty_cache() |
|
column_names = [f"{self.model_name}_feat_{ii}" for ii in range(len(classification_outputs_tensor[0]))] |
|
return pd.DataFrame( |
|
data=classification_outputs_tensor.tolist(), |
|
index=data.index, |
|
columns=column_names |
|
) |
|
|
|
|
|
class ManyBertPretrainFeatureExtractor(BaseExtractor): |
|
def __init__(self, model_names: List[str], max_length: int = 512, batch_size=64): |
|
super(ManyBertPretrainFeatureExtractor, self).__init__() |
|
self.model_names = model_names |
|
self.max_length = max_length |
|
self.batch_size = batch_size |
|
|
|
def generate_features(self, X: pd.Series) -> pd.DataFrame: |
|
extractors = [ |
|
BertPretrainFeatureExtractor(model_name, self.max_length, self.batch_size) |
|
for model_name in self.model_names |
|
] |
|
dataframes = [ |
|
extractor.generate_features(X) for extractor in extractors |
|
] |
|
|
|
return pd.concat(dataframes, axis='columns') |
|
|
