utils2.py

import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from tqdm import *
import numpy as np
import logging as log
from collections import Counter,defaultdict
import os
import time
from datasets import load_dataset
from numba import jit, prange, njit, guvectorize
from simsimd import cdist, DistancesTensor, cosine
from sentence_transformers import util
import heapq
import lzma
import pickle
import spacy
import json
import zstandard as zstd
import pickle
import zipfile

from nltk.corpus import wordnet as wn
from collections import defaultdict

WN_POS_MAP = {'n': 'NOUN', 'v': 'VERB', 'a': 'ADJ', 'r': 'ADV', 's': 'ADJ'}

def load_zst(path):
    dctx = zstd.ZstdDecompressor()
    with open(path, "rb") as f:
        return pickle.loads(dctx.decompress(f.read()))

def get_pos_tags(nlp, sentence):
    sentence_doc = nlp(sentence)
    split_to_spacy = {}
    char_pos = 0
    for idx, word in enumerate(sentence.split()):
        start = sentence.index(word, char_pos)
        for token in sentence_doc:
            if token.idx == start:
                split_to_spacy[idx] = token
                break
        char_pos = start + len(word)
    return split_to_spacy        

def sample_santext_gpu_ED_pos(embeddings, word_emb, epsilon, pos_mask):
    dists = torch.cdist(word_emb.unsqueeze(0), embeddings, p=2)[0]
    logits = -epsilon * dists / 2
    logits[~pos_mask] = float('-inf')
    probs = torch.softmax(logits, dim=0)
    idx = torch.multinomial(probs, 1).item()
    return idx
    
def get_wordnet_pos(word):
    synsets = wn.synsets(word)
    return set(WN_POS_MAP[s.pos()] for s in synsets)

def discover_corpus_pos(nlp, corpus, batch_size=1000, max_sentences=50000):
    word_pos = defaultdict(set)
    sentences = corpus[:max_sentences]
    for i in trange(0, len(sentences), batch_size, desc="Corpus PoS"):
        batch = sentences[i:i + batch_size]
        for doc in nlp.pipe(batch):
            for token in doc:
                word_pos[token.text].add(token.pos_)
    return word_pos

def build_pos_maps(nlp, idx2word, corpus=None):
    corpus_pos = defaultdict(set)
    if corpus:
        batch_size = 1000
        sentences = corpus
        for i in trange(0, len(sentences), batch_size, desc="Corpus PoS discovery"):
            batch = sentences[i:i + batch_size]
            for doc in nlp.pipe(batch):
                for token in doc:
                    corpus_pos[token.text].add(token.pos_)

    word2pos = {}
    batch_size = 5000
    for i in trange(0, len(idx2word), batch_size, desc="Building PoS maps"):
        batch = [str(w) for w in idx2word[i:i + batch_size]]
        docs = list(nlp.pipe(batch))
        for w, doc in zip(batch, docs):
            wn_pos = get_wordnet_pos(w)
            tokens = list(doc)
            spacy_pos = {tokens[0].pos_} if tokens else set()
            combined = wn_pos | spacy_pos | corpus_pos.get(w, set())
            if not combined:
                combined = {"X"}
            word2pos[w] = combined
            
    # pos_array = [sorted(word2pos.get(str(w), {"X"})) for w in idx2word]

    pos_array = [word2pos.get(str(w), {"X"}) for w in idx2word]
    return word2pos, pos_array


def get_possible_pos_word2net(word, wn = None):
    synsets = wn.synsets(word)
    pos_map = {'n': 'NOUN', 'v': 'VERB', 'a': 'ADJ', 'r': 'ADV', 's': 'ADJ'}
    return set(pos_map[s.pos()] for s in synsets)

#print(get_possible_pos("light"))  # {'NOUN', 'VERB', 'ADJ', 'ADV'}

def get_contextualized_pos(nlp, sentence, word, position=None):
    doc = nlp(sentence)
    matches = [(token.i, token.pos_) for token in doc if token.text == word]
    if not matches:
        return None
    if position is not None:
        for idx, pos in matches:
            if idx == position:
                return pos
        return None
    return matches[0][1]

def build_pos_tags(idx2word, nlp = None):
        """Pre-compute PoS tags for all words in the embedding vocabulary.
        
        Uses spaCy's single-token PoS tagger. Since we're tagging isolated words
        (no sentence context), this gives a "most likely" PoS per word form.
        Returns a numpy array aligned with idx2word for fast indexing.
        """
        if nlp == None:
            nlp = spacy.load("en_core_web_sm", disable=["ner", "lemmatizer"])
        print("Pre-computing PoS tags for vocabulary...")
        word2pos = {}
        batch_size = 5000
        for i in trange(0, len(idx2word), batch_size, desc="PoS tagging"):
            batch = [str(w) for w in idx2word[i:i + batch_size]]
            docs = nlp.pipe(batch)
            for word, doc in zip(batch, docs):
                tokens = list(doc)
                if len(tokens) > 0:
                    word2pos[word] = tokens[0].pos_
                else:
                    word2pos[word] = "X"  # unknown
        
        # Build a PoS array aligned with idx2word for fast lookup by index
        pos_array = np.array([word2pos.get(w, "X") for w in idx2word])
        print(f"PoS tag distribution: {Counter(pos_array).most_common(15)}")
        return word2pos, pos_array
        
def select_topk_by_pos_gpu(word_idx, target_pos, embeddings_gpu, pos_array, device, topk=5, pos_masks=None):
    if pos_masks is not None:
        pos_mask = pos_masks.get(target_pos, torch.zeros(len(pos_array), dtype=torch.bool, device=device))
    else:
        pos_mask = torch.tensor(
            [target_pos in pos_array[j] for j in range(len(pos_array))],
            device=device
        )

    word_emb = embeddings_gpu[word_idx].unsqueeze(0)
    all_scores = torch.cdist(word_emb, embeddings_gpu, p=2)[0]
    all_scores[~pos_mask] = float('inf')
    all_scores[all_scores >= 1e8] = float('inf')

    k = min(topk, pos_mask.sum().item())
    top_indices = torch.tensor([], dtype=torch.long, device=device)
    if k > 0:
        _, top_indices = torch.topk(all_scores, k, largest=False)

    if len(top_indices) < topk:
        needed = topk - len(top_indices)
        fb_scores = torch.cdist(word_emb, embeddings_gpu, p=2)[0]
        fb_scores[fb_scores >= 1e8] = float('inf')
        if len(top_indices) > 0:
            fb_scores[top_indices] = float('inf')
        k_fb = min(needed, (fb_scores != float('inf')).sum().item())
        if k_fb > 0:
            _, fb_idx = torch.topk(fb_scores, k_fb, largest=False)
            top_indices = torch.cat([top_indices, fb_idx])

    return top_indices.cpu().numpy()

def select_topk_pos_from_sentence(nlp, sentence, position, word2idx, embeddings_gpu, pos_array, device, topk=20):
    if not isinstance(sentence, list):
        words = sentence.split()
    word = words[position]

    if word not in word2idx:
        return None

    doc = nlp(sentence)
    # Align split position to spaCy token
    char_pos = sum(len(w) + 1 for w in words[:position])
    target_pos = "X"
    for token in doc:
        if token.idx == char_pos:
            target_pos = token.pos_
            break

    word_idx = word2idx[word]
    return select_topk_by_pos_gpu(word_idx, target_pos, embeddings_gpu, pos_array, device, topk)

def find_matching_positions(values, searchval):
    N = len(searchval)
    possibles = np.where(values == searchval[0])[0]

    solns = []
    for p in possibles:
        check = values[p:p+N]
        if np.all(check == searchval):
            return p
    return None


def search(pattern, text, q):
    m = len(pattern)
    n = len(text)
    p = 0
    t = 0
    h = 1
    i = 0
    j = 0

    for i in range(m-1):
        h = (h*d) % q

    # Calculate hash value for pattern and text
    for i in range(m):
        p = (d*p + ord(pattern[i])) % q
        t = (d*t + ord(text[i])) % q

    # Find the match
    for i in range(n-m+1):
        if p == t:
            for j in range(m):
                if text[i+j] != pattern[j]:
                    break

            j += 1
            if j == m:
                print("Pattern is found at position: " + str(i+1))

        if i < n-m:
            t = (d*(t-ord(text[i])*h) + ord(text[i+m])) % q

            if t < 0:
                t = t+q

@jit(nopython=True)
def get_most_similar_token_topk(word_embeddings, new_embedd_vector, skip, topk):
    return fast_cosine_matrix(new_embedd_vector, word_embeddings).argsort()[::-1][1:topk]
    scores = fast_cosine_matrix(new_embedd_vector, word_embeddings) #.argsort()[::-1][:topk]
    #print(scores)
    #indices = fast_cosine_matrix(new_embedd_vector, word_embeddings)#.argsort()[::-1][:topk]
    #exit(1)
    return scores#, indices
    #for i, emb_vector in enumerate(word_embeddings):
    #    scores.append(fast_cosine_matrix(emb_vector, new_embedd_vector))
    
    for i in range(len(scores)):
        if i == skip:
            continue
        if scores[i] > max_score:
            max_score = scores[i]
            nearest_token = i
    return nearest_token

@jit(nopython=True)
def get_most_similar_token(word_embeddings, new_embedd_vector, skip):
    scores = fast_cosine_matrix(new_embedd_vector, word_embeddings)

    scores[skip] = -np.inf  # Exclude the skipped index
    nearest_token = np.argmax(scores)
    return nearest_token

def get_most_similar_token_faster(word_embeddings, new_embedd_vector, skip):
    if word_embeddings.dtype != new_embedd_vector.dtype:
        new_embedd_vector = np.array(new_embedd_vector, dtype=word_embeddings.dtype)
    
    # distance = cosine(new_embedd_vector, word_embeddings)
    # score = np.array(distance, copy=True) 
    # score = 1 - score
    # score[skip] = -np.inf
    # return int(np.argmax(score))
    distance = cdist(new_embedd_vector, word_embeddings, metric='cosine', threads=0)
    score = np.array(distance, copy=True) 
    score = score[0]
    score = 1 - score
    if skip != -1:
        score[skip] = -np.inf
    return int(np.argmax(score))

def get_most_similar_token_gpu(word_embeddings, new_embedd_vector, skip):
    with torch.no_grad():
        if word_embeddings.dtype != new_embedd_vector.dtype:
            new_embedd_vector = np.array(new_embedd_vector, dtype=word_embeddings.dtype)
        
        #new_embedd_vector = torch.from_numpy(new_embedd_vector)
        cosine_similarity = util.pytorch_cos_sim(word_embeddings, new_embedd_vector)
        cosine_distance = 1 - cosine_similarity
        if skip != -1:
            cosine_distance[skip] = np.inf
        a = torch.argmin(cosine_distance, dim=0)
        return int(a.cpu())
    #return int(np.argmin(cosine_distance.cpu()))

def get_most_similar_token_gpu_vector(word_embeddings, new_embedd_vector, skip):

    if word_embeddings.dtype != new_embedd_vector.dtype:
        new_embedd_vector = np.array(new_embedd_vector, dtype=word_embeddings.dtype)
    cosine_similarity = util.pytorch_cos_sim(word_embeddings, new_embedd_vector)
    cosine_distance = 1 - cosine_similarity
    if skip != -1:
        cosine_distance[skip] = np.inf
    return cosine_distance

def get_most_similar_token_gpuED(word_embeddings, new_embedd_vector, skip):
    with torch.no_grad():
        if word_embeddings.dtype != new_embedd_vector.dtype:
            new_embedd_vector = np.array(new_embedd_vector, dtype=word_embeddings.dtype)
        ed = torch.cdist(new_embedd_vector.unsqueeze(0), word_embeddings, p = 2)[0]
        if skip != -1:
            ed[skip] = np.inf
    return  int(torch.argmin(ed).cpu())

def get_most_similar_token_gpu_topkED(word_embeddings, new_embedd_vector, skip, topk):
    with torch.no_grad():
        if word_embeddings.dtype != new_embedd_vector.dtype:
            new_embedd_vector = np.array(new_embedd_vector, dtype=word_embeddings.dtype)
        ed = torch.cdist(new_embedd_vector.unsqueeze(0), word_embeddings, p = 2)[0]
        if skip != -1:
            ed[skip] = np.inf
        sorted_indices = torch.argsort(ed, dim=0)
        top_5_smallest_indices = sorted_indices[:5]
    return top_5_smallest_indices.cpu()

def get_most_similar_token_gpu_topkworks(word_embeddings, new_embedd_vector, topk):
    with torch.no_grad():
        if word_embeddings.dtype != new_embedd_vector.dtype:
            new_embedd_vector = np.array(new_embedd_vector, dtype=word_embeddings.dtype)
        ed = torch.cdist(new_embedd_vector.unsqueeze(0), word_embeddings, p = 2)[0]
        sorted_indices = torch.argsort(ed, dim=0)
        top_5_smallest_indices = sorted_indices[:topk]
    return top_5_smallest_indices.cpu()

def sample_santext_gpu_ED(word_embeddings, new_embedd_vector, epsilon):
    with torch.no_grad():
        if word_embeddings.dtype != new_embedd_vector.dtype:
            new_embedd_vector = new_embedd_vector.to(word_embeddings.dtype)
        ed = torch.cdist(new_embedd_vector.unsqueeze(0), word_embeddings, p=2)[0]
                
        logits = -epsilon * ed / 2
        full_probs = torch.softmax(logits, dim=0).cpu().numpy()
        sampled_index = np.random.choice(len(full_probs), 1, p=full_probs)[0]
        # DEBUG - remove after

        return sampled_index
    
def get_most_similar_token_gpu_topk(word_embeddings, new_embedd_vector, skip, topk):

    if word_embeddings.dtype != new_embedd_vector.dtype:
        new_embedd_vector = np.array(new_embedd_vector, dtype=word_embeddings.dtype)

    cosine_similarity = util.pytorch_cos_sim(word_embeddings, new_embedd_vector)
    cosine_distance = 1 - cosine_similarity
    if skip != -1:
        cosine_distance[skip] = np.inf
    sorted_indices = torch.argsort(cosine_distance, dim=0)
    top_5_smallest_indices = sorted_indices[:5]
    return top_5_smallest_indices.cpu()

#@jit(nopython=True)
def get_most_similar_tokens(word_embeddings, new_embedd_vectors, skips):
    t1 = time.time()
    nearest_token = np.zeros(new_embedd_vectors.shape[0])
    for i, vector in enumerate(new_embedd_vectors):
        distance = cosine(vector, word_embeddings)
        score = np.array(distance, copy=True) 

        score[skips[i]] = np.inf
        nearest_token[i] = int(np.argmin(score))
    print('Time first:', time.time() - t1)
    print(nearest_token)
    # for i in range(new_embedd_vectors.shape[0]):
    #    #distances = cdist(new_embedd_vectors[i], word_embeddings, metric="cosine")
    #    #scores = np.array(distances, copy=True)  
    #    scores = fast_cosine_matrix(new_embedd_vectors[i], word_embeddings)
    #    scores[skips[i]] = -np.inf  # Exclude the skipped index
    #    nearest_token[i] = int(np.argmax(scores))
    return nearest_token

@guvectorize(["void(float64[:], float64[:], float64[:])"], "(n),(n)->()", target='parallel')
def fast_cosine_gufunc(u, v, result):
    m = u.shape[0]
    udotv = 0
    u_norm = 0
    v_norm = 0
    for i in range(m):
        if (np.isnan(u[i])) or (np.isnan(v[i])):
            continue

        udotv += u[i] * v[i]
        u_norm += u[i] * u[i]
        v_norm += v[i] * v[i]

    u_norm = np.sqrt(u_norm)
    v_norm = np.sqrt(v_norm)

    if (u_norm == 0) or (v_norm == 0):
        ratio = 1.0
    else:
        ratio = udotv / (u_norm * v_norm)
    result[:] = ratio

@jit(nopython=True, parallel=True)
def fast_cosine_matrix(u, M):
    scores = np.zeros(M.shape[0])
    for i in prange(M.shape[0]):
        v = M[i]
        m = u.shape[0]
        udotv = 0
        u_norm = 0
        v_norm = 0
        for j in range(m):
            udotv += u[j] * v[j]
            u_norm += u[j] * u[j]
            v_norm += v[j] * v[j]

        u_norm = np.sqrt(u_norm)
        v_norm = np.sqrt(v_norm)

        if (u_norm == 0) or (v_norm == 0):
            ratio = 1.0
        else:
            ratio = udotv / (u_norm * v_norm)
        scores[i] = ratio
    return scores

@jit(nopython=True, parallel=True)
def fast_cosine_matrices(U, M):
    """
    Calculate the cosine similarity between two matrices U and M.

    Args:
        U (numpy.ndarray): A matrix of shape (n_rows_U, n_features).
        M (numpy.ndarray): A matrix of shape (n_rows_M, n_features).

    Returns:
        numpy.ndarray: A 2D array of shape (n_rows_U, n_rows_M), where the element at
                       (i, j) is the cosine similarity between U[i] and M[j].
    """
    n_rows_U = U.shape[0]
    n_rows_M = M.shape[0]
    n_features = U.shape[1]

    scores = np.zeros((n_rows_U, n_rows_M))

    for i in prange(n_rows_U):
        for j in range(n_rows_M):
            udotv = 0.0
            u_norm = 0.0
            v_norm = 0.0

            for k in range(n_features):
                udotv += U[i, k] * M[j, k]
                u_norm += U[i, k] * U[i, k]
                v_norm += M[j, k] * M[j, k]

            u_norm = np.sqrt(u_norm)
            v_norm = np.sqrt(v_norm)

            if (u_norm == 0) or (v_norm == 0):
                scores[i, j] = 1.0  # Assuming similarity of 1.0 for degenerate cases
            else:
                scores[i, j] = udotv / (u_norm * v_norm)

    return scores

def get_text_from_input_ids(tokenizer, input_ids, skip_special=False):
    """
    from input ids it returns the input ids
    """
    texts = []
    if not isinstance(input_ids[0], list):
        return tokenizer.decode(input_ids, clean_up_tokenization_spaces = True, skip_special_tokens = skip_special)
    for ids in input_ids:
        text = tokenizer.decode(ids, clean_up_tokenization_spaces = True, skip_special_tokens = skip_special)
        texts.append(text)
    return texts


def get_input_ids_from_text(tokenizer, data, with_special_tokens=True):
    input_ids = []
    if isinstance(data[0], list):
        for sent in tqdm(data):
            text = tokenizer.encode(sent, add_special_tokens=with_special_tokens, pad_to_max_length = not with_special_tokens)  # Pad sentence to max lengthtruncation=not with_special_tokens)
            input_ids.append(text)
    else:
        return tokenizer.encode(data, add_special_tokens=with_special_tokens, pad_to_max_length = not with_special_tokens) 
    #log.info("Done generating input ids from text!")
    return input_ids



def encode_text(tokenizer, data, max_len, with_seperation=True):
    # Create empty lists to store outputs
    input_ids = []
    attention_masks = []

    # For every sentence...
    for sent in data:
        # sent = sent["sentence"]
        # `encode_plus` will:
        #    (1) Tokenize the sentence
        #    (2) Add the `[CLS]` and `[SEP]` token to the start and end
        #    (3) Truncate/Pad sentence to max length
        #    (4) Map tokens to their IDs
        #    (5) Create attention mask
        #    (6) Return a dictionary of outputs
        encoded_sent = tokenizer.encode_plus(
            text = sent,  # Preprocess sentence
            add_special_tokens = with_seperation,  # Add `[CLS]` and `[SEP]`
            max_length = max_len,  # Max length to truncate/pad
            pad_to_max_length = with_seperation,  # Pad sentence to max length
            # return_tensors='pt',           # Return PyTorch tensor
            return_attention_mask = True  # Return attention mask
        )
        # Add the outputs to the lists
        input_ids.append(encoded_sent.get('input_ids'))
        attention_masks.append(encoded_sent.get('attention_mask'))

    # Convert lists to tensors
    if with_seperation:
        input_ids = torch.tensor(input_ids)
        attention_masks = torch.tensor(attention_masks)

    return input_ids, attention_masks

# def create_dataloader_roberta(train_dataset, val_dataset, train_labels, val_labels, batch_size):
#     train_data = TensorDataset(train_input_ids, train_attention_mask, train_labels)

#     train_dataloader = DataLoader(train_dataset, batch_size = batch_size, drop_last=True, shuffle=True, num_workers=0)
#     val_dataloader = DataLoader(val_dataset, batch_size=batch_size, num_workers=0)
#     return train_dataloader, val_dataloader

# create dataloader for training
def create_dataloader(train_input_ids, train_attention_mask, train_labels, batch_size):
    train_data = TensorDataset(train_input_ids, train_attention_mask, train_labels)
    train_sampler = RandomSampler(train_data)
    train_dataloader = DataLoader(train_data, sampler = train_sampler, batch_size = batch_size)

    return train_data, train_sampler, train_dataloader


# get the maximum length of a word
def get_max_len(dataset, tokenizer):
    sentence_train = [sent["sentence"] for sent in dataset["train"]]
    sentence_validation = [sent["sentence"] for sent in dataset["validation"]]
    sentence_test = [sent["sentence"] for sent in dataset["test"]]

    data = np.concatenate([sentence_train, sentence_validation, sentence_test])

    # Encode our concatenated data
    encoded_data = [tokenizer.encode(sent, add_special_tokens = True) for sent in data]

    # Find the maximum length
    max_len = max([len(sent) for sent in encoded_data])
    print('Max length: ', max_len)
    return max_len

# create dictonairy with text and label from sst2-like datasets
def create_data(dataset):
    train_val_dict = []
    data = load_dataset(dataset)#, split='unsupervised')

    if "sst2" in dataset:
        text_key = "sentence"
        types = ["train", "validation"]
    elif "imdb" in dataset:
        text_key = "text"
        types = ["train", "test"]

    for type in types:
        df = {}
        texts = [sent[text_key] for sent in data[type]]
        labels = [sent["label"] for sent in data[type]]
        df["text"] = texts
        df["label"] = labels
        train_val_dict.append(df)
    return train_val_dict[0], train_val_dict[1]


def get_input_ids_frequency(tokenizer, input_ids):
    vocab = text_manipulation.create_vocabulary(tokenizer=tokenizer)

    freq_ids_map = {int(key): 0 for key in vocab}
    for i, ids in enumerate(tqdm(input_ids)):
        for val in ids:
            val = int(val)
            if val == 0:
                continue
            freq_ids_map[val] = freq_ids_map[val] + 1
    log.info("Done generating input ids from text!")
    return freq_ids_map


def create_vocabulary(tokenizer, vocab_type="token"):
    if vocab_type == "token":
        d = tokenizer.get_vocab()
        arr = [0] * len(d)
        for i, key in enumerate(d):
            arr[i] = d[key]
        return arr

    if vocab_type == "word":
        exit(1)
        word_freqs = defaultdict(int)
        pre_tokenizer = tokenizers.pre_tokenizers.BertPreTokenizer()
        for i, sent in enumerate(data):
            words_with_offsets = pre_tokenizer.pre_tokenize_str(sent)
            new_words = [word for word, offset in words_with_offsets]
            for word in new_words:
                if word in string.punctuation:
                    continue
                word_freqs[str(word)] += 1
        return word_freqs


def gaussian_weights(size, sigma=1):
    x = np.linspace(-(size // 2), size // 2, size)
    gaussian = np.exp(-x**2 / (2 * sigma**2))
    return gaussian / np.sum(gaussian)  # Normalize the weights to sum up to 1


@njit(fastmath=True,parallel=True)
def eucl_naive(A, B):
    assert A.shape[1]==B.shape[1]
    C=np.empty((A.shape[0],B.shape[0]),A.dtype)
    
    #workaround to get the right datatype for acc
    init_val_arr=np.zeros(1,A.dtype)
    init_val=init_val_arr[0]
    
    for i in prange(A.shape[0]):
        for j in range(B.shape[0]):
            acc=init_val
            for k in range(A.shape[1]):
                acc+=(A[i,k]-B[j,k])**2
            C[i,j]=np.sqrt(acc)
    return C

#@jit(fastmath=True,parallel=True)
def generate_points_within_hypersphere(num_points, num_dimensions, radius=1):
    points = []
    for _ in range(num_points):
        point = np.random.normal(size=num_dimensions)
        point /= np.linalg.norm(point)  # Normalize to unit vector
        point *= radius  # Scale to desired radius
        points.append(point)
    return np.array(points)

def load(save_path, compress=False):
        if not os.path.exists(save_path):
            if ".json" in save_path:
                if "token_has" in save_path:
                    return defaultdict(int)
                elif "word_has" in save_path:
                    return defaultdict(str)
                else:
                    return defaultdict(list)
            elif ".tsv" in save_path:
                return []
            elif ".txt" in save_path:
                return []
            else:
                return []
        else:
            if os.path.exists(save_path) and os.stat(save_path).st_size > 1:
                if compress:
                    with lzma.open(save_path, "rb") as f:
                        return pickle.load(f)
                else:
                    with open(save_path, 'r') as f:
                        if ".json" in save_path:
                            return json.load(f)
                        elif ".tsv" in save_path:
                            return pd.read_csv(f, sep='\t')
                        elif ".txt" in save_path:
                            return f.readlines()


def checkpoint2(save_path, data, compress=False):
    if compress:
        with lzma.open(save_path, "wb") as f:
            pickle.dump(data, f)
    else:
        with open(save_path, 'w') as f:
            if ".json" in save_path:
                f.write(json.dumps(data, ensure_ascii=False, indent=4))
            elif ".tsv" in save_path:
                data.to_csv(save_path, sep="\t", index=0)
            elif ".txt" in save_path:
                f.writelines()

def is_number(s):
        try:
            float(s)
            return True
        except ValueError:
            pass
    
        try:
            import unicodedata
            unicodedata.numeric(s)
            return True
        except (TypeError, ValueError):
            pass
    
        return False

def create_glove_embeddings_table(out_prefix_path, glove_path):
    word2idx = []
    idx2word = []
    with open(glove_path,'r') as file:
        for i,line in tqdm(enumerate(file)):
            stripped = line.strip().split()
            valuess = line.strip().split()[-300:]
            j = len(stripped) - 300
            embedding = [float(num) for num in valuess]
            embeddings.append(embedding)
            idx2word.append("".join(stripped[0:j]))
            word2idx["".join(stripped[0:j])] = i
    embeddings = np.array(embeddings)
    idx2word = np.asarray(idx2word)
    print('calculating the linalg')
    norm = np.linalg.norm(embeddings, axis=1, keepdims=True)
    embeddings = np.asarray(embeddings / norm, "float64")
    print('done.. saving')
    checkpoint2("emb" + out_prefix_path + '.pkl.xz', embeddings, True)
    checkpoint2("idx2word" + out_prefix_path + '.pkl.xz', embeddings, True)
    checkpoint2("word2idx" + out_prefix_path + '.pkl.xz', embeddings, True)
    #self.checkpoint(idx2word_save_path, idx2word, True)
    #self.checkpoint(word2idx_save_path, word2idx, True)