File size: 2,995 Bytes
8b414b0 |
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 |
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')
|