|
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) |
|
|
|
|
|
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 |
|
""" |
|
|
|
|
|
|
|
with torch.no_grad(): |
|
outputs = model(tokens_tensor, segments_tensors) |
|
|
|
|
|
hidden_states = outputs[2][1:] |
|
|
|
|
|
token_embeddings = hidden_states[-1] |
|
|
|
token_embeddings = torch.squeeze(token_embeddings, dim=0) |
|
|
|
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 = [] |
|
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 |
|
|
|
|
