import pandas as pd import numpy as np import pickle import torch from torch.utils.data import Dataset, DataLoader from transformers import BertTokenizer, BertModel from transformers import AutoTokenizer, AutoModel import nltk tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') model = BertModel.from_pretrained('bert-base-uncased', output_hidden_states = True,) def extract_context_words(x, window = 6): paragraph, offset_start, offset_end = x['paragraph'], x['offset_start'], x['offset_end'] target_word = paragraph[offset_start : offset_end] paragraph = ' ' + paragraph + ' ' offset_start = offset_start + 1 offset_end = offset_end + 1 prev_space_posn = (paragraph[:offset_start].rindex(' ') + 1) end_space_posn = (offset_end + paragraph[offset_end:].index(' ')) full_word = paragraph[prev_space_posn : end_space_posn] prev_words = nltk.word_tokenize(paragraph[0:prev_space_posn]) next_words = nltk.word_tokenize(paragraph[end_space_posn:]) words_in_context_window = prev_words[-1*window:] + [full_word] + next_words[:window] context_text = ' '.join(words_in_context_window) return context_text """The following functions have been created with inspiration from https://github.com/arushiprakash/MachineLearning/blob/main/BERT%20Word%20Embeddings.ipynb""" def bert_text_preparation(text, tokenizer): """Preparing the input for BERT Takes a string argument and performs pre-processing like adding special tokens, tokenization, tokens to ids, and tokens to segment ids. All tokens are mapped to seg- ment id = 1. Args: text (str): Text to be converted tokenizer (obj): Tokenizer object to convert text into BERT-re- adable tokens and ids Returns: list: List of BERT-readable tokens obj: Torch tensor with token ids obj: Torch tensor segment ids """ marked_text = "[CLS] " + text + " [SEP]" tokenized_text = tokenizer.tokenize(marked_text) indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text) segments_ids = [1]*len(indexed_tokens) # Convert inputs to PyTorch tensors tokens_tensor = torch.tensor([indexed_tokens]) segments_tensors = torch.tensor([segments_ids]) return tokenized_text, tokens_tensor, segments_tensors def get_bert_embeddings(tokens_tensor, segments_tensors, model): """Get embeddings from an embedding model Args: tokens_tensor (obj): Torch tensor size [n_tokens] with token ids for each token in text segments_tensors (obj): Torch tensor size [n_tokens] with segment ids for each token in text model (obj): Embedding model to generate embeddings from token and segment ids Returns: list: List of list of floats of size [n_tokens, n_embedding_dimensions] containing embeddings for each token """ # Gradient calculation id disabled # Model is in inference mode with torch.no_grad(): outputs = model(tokens_tensor, segments_tensors) # Removing the first hidden state # The first state is the input state hidden_states = outputs[2][1:] # Getting embeddings from the final BERT layer token_embeddings = hidden_states[-1] # Collapsing the tensor into 1-dimension token_embeddings = torch.squeeze(token_embeddings, dim=0) # Converting torchtensors to lists list_token_embeddings = [token_embed.tolist() for token_embed in token_embeddings] return list_token_embeddings def bert_embedding_extract(context_text, word): tokenized_text, tokens_tensor, segments_tensors = bert_text_preparation(context_text, tokenizer) list_token_embeddings = get_bert_embeddings(tokens_tensor, segments_tensors, model) word_tokens,tt,st = bert_text_preparation(word, tokenizer) word_embedding_all = [] try: for word_tk in word_tokens: word_index = tokenized_text.index(word_tk) word_embedding = list_token_embeddings[word_index] word_embedding_all.append(word_embedding) word_embedding_mean = np.array(word_embedding_all).mean(axis=0) return word_embedding_mean except: return ['None']