|
from transformers import PreTrainedModel |
|
from nltk.corpus import stopwords |
|
from nltk.tokenize import word_tokenize |
|
from .config_gzipembed import * |
|
from tqdm.auto import tqdm |
|
import torch |
|
import gzip |
|
import multiprocessing |
|
class GZIPEmbeddingModel(PreTrainedModel): |
|
config_class = GZIPEmbeddingConfig |
|
def __init__(self, config): |
|
super().__init__(config) |
|
if config.reduction: |
|
self.reduction_head = torch.nn.Linear(len(config.corpus), config.reduced_dimension) |
|
else: |
|
self.reduction_head = None |
|
self.dummy_parameter = torch.nn.Parameter(torch.ones(1)) |
|
|
|
def forward(self, prompt, num_procs=16, return_tensor=True): |
|
global calculate_ncd_row |
|
global p |
|
def calculate_ncd_row(data_row): |
|
i = data_row[0] |
|
row = self.ncd(data_row[1], p) |
|
return i, row |
|
if type(prompt) == str: |
|
prompt = [prompt] |
|
x = [] |
|
for p in prompt: |
|
ncd = [0] * len(self.config.corpus) |
|
with multiprocessing.Pool(num_procs) as pool: |
|
data = enumerate(self.config.corpus) |
|
results = pool.map(calculate_ncd_row,data) |
|
for i,row in results: |
|
ncd[i]=row |
|
x.append(ncd) |
|
if self.reduction_head is not None: |
|
x = torch.tensor(x) |
|
x = x.to(self.reduction_head.dtype).to(self.reduction_head.device) |
|
return self.reduction_head(x) |
|
return x if not return_tensor else torch.tensor(x) |
|
|
|
def encode(self, sentences, batch_size=32, **kwargs): |
|
""" |
|
Returns a list of embeddings for the given sentences. |
|
Args: |
|
sentences (`List[str]`): List of sentences to encode |
|
batch_size (`int`): Batch size for the encoding |
|
|
|
Returns: |
|
`List[np.ndarray]` or `List[tensor]`: List of embeddings for the given sentences |
|
""" |
|
import numpy as np |
|
x = self.forward(sentences, num_procs=batch_size, return_tensor=False) |
|
|
|
return [np.array(i) for i in x] |
|
|
|
def normalize(self, x): |
|
x = ''.join([char for char in x.lower() if char in "abcdefghijklmnopqrstuvwxyz "]) |
|
x = word_tokenize(x) |
|
x = [w for w in x if not w in self.config.stop_words] |
|
return ' '.join(x) |
|
|
|
def ncd(self, x, y): |
|
_x = self.normalize(x) if self.config.normalize else x |
|
_y = self.normalize(y) if (not self.config.normalized_corpus) and self.config.normalize else y |
|
x_c = len(gzip.compress(_x.encode())) |
|
y_c = len(gzip.compress(_y.encode())) |
|
xy_c = len(gzip.compress(f"{_x} {_y}".encode())) |
|
return (xy_c-min(x_c,y_c))/max(x_c,y_c) |
|
|
|
def gzip_embed( |
|
self, |
|
corpus, |
|
document, |
|
verbose=False, |
|
): |
|
embedding = [] |
|
for reference_document in (corpus if not verbose else tqdm(corpus)): |
|
embedding.append(self.ncd(reference_document, document)) |
|
return embedding |
|
|
|
def dimensionality(self): |
|
return len(self.config.corpus) |
|
|
