Upload fincat_utils.py

fincat_utils.py

@@ -0,0 +1,108 @@
+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 = []
+  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
+