notebooks/BERTembeddings.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### generate word embeddings with BERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertModel, BertTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BertModel.from_pretrained('bert-base-uncased')\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['i', 'really', 'like', 'this', 'example', 'sentence']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"i really like this example sentence\"\n",
    "tokens = tokenizer.tokenize(text)\n",
    "tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1045, 2428, 2066, 2023, 2742, 6251]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "input_ids = torch.tensor(input_ids).unsqueeze(0)\n",
    "with torch.no_grad():\n",
    "    outputs = model(input_ids)\n",
    "    embeddings = outputs.last_hidden_state[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### word embeddings with BERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.1693,  0.2022,  0.1156,  ...,  0.0366,  0.2253,  0.1816],\n",
       "        [ 0.6730,  0.6221,  0.3550,  ..., -0.1383, -0.0113,  0.0695],\n",
       "        [ 0.3303,  0.6492,  0.3663,  ..., -0.3147,  0.0859, -0.0903],\n",
       "        [-0.3525,  0.3459, -0.0413,  ..., -0.2913,  0.1090,  0.2372],\n",
       "        [-0.1793,  0.5224, -0.0567,  ..., -0.1681,  0.1074,  0.0227],\n",
       "        [-0.1854,  0.5431, -0.0495,  ..., -0.1403,  0.1456, -0.0548]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertModel, BertTokenizer\n",
    "model = BertModel.from_pretrained('bert-base-uncased',\n",
    "           output_hidden_states = True,)\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bert_text_preparation(text, tokenizer):\n",
    "    \"\"\"\n",
    "    Preprocesses text input in a way that BERT can interpret.\n",
    "    \"\"\"\n",
    "    marked_text = \"[CLS] \" + text + \" [SEP]\"\n",
    "    tokenized_text = tokenizer.tokenize(marked_text)\n",
    "    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)\n",
    "    segments_ids = [1]*len(indexed_tokens)\n",
    "#   convert inputs to tensors\n",
    "    tokens_tensor = torch.tensor([indexed_tokens])\n",
    "    segments_tensor = torch.tensor([segments_ids])\n",
    "    return tokenized_text, tokens_tensor, segments_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_bert_embeddings(tokens_tensor, segments_tensor, model):\n",
    "    \"\"\"\n",
    "    Obtains BERT embeddings for tokens.\n",
    "    \"\"\"\n",
    "    # gradient calculation id disabled\n",
    "    with torch.no_grad():\n",
    "      # obtain hidden states\n",
    "      outputs = model(tokens_tensor, segments_tensor)\n",
    "      hidden_states = outputs[2]\n",
    "    # concatenate the tensors for all layers\n",
    "    # use \"stack\" to create new dimension in tensor\n",
    "    token_embeddings = torch.stack(hidden_states, dim=0)\n",
    "    # remove dimension 1, the \"batches\"\n",
    "    token_embeddings = torch.squeeze(token_embeddings, dim=1)\n",
    "    # swap dimensions 0 and 1 so we can loop over tokens\n",
    "    token_embeddings = token_embeddings.permute(1,0,2)\n",
    "    # intialized list to store embeddings\n",
    "    token_vecs_sum = []\n",
    "    # \"token_embeddings\" is a [Y x 12 x 768] tensor\n",
    "    # where Y is the number of tokens in the sentence\n",
    "    # loop over tokens in sentence\n",
    "    for token in token_embeddings:\n",
    "    # \"token\" is a [12 x 768] tensor\n",
    "    # sum the vectors from the last four layers\n",
    "        sum_vec = torch.sum(token[-4:], dim=0)\n",
    "        token_vecs_sum.append(sum_vec)\n",
    "    return token_vecs_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentences = [\"bank\",\n",
    "         \"he eventually sold the shares back to the bank at a premium.\",\n",
    "         \"the bank strongly resisted cutting interest rates.\",\n",
    "         \"the bank will supply and buy back foreign currency.\",\n",
    "         \"the bank is pressing us for repayment of the loan.\",\n",
    "         \"the bank left its lending rates unchanged.\",\n",
    "         \"the river flowed over the bank.\",\n",
    "         \"tall, luxuriant plants grew along the river bank.\",\n",
    "         \"his soldiers were arrayed along the river bank.\",\n",
    "         \"wild flowers adorned the river bank.\",\n",
    "         \"two fox cubs romped playfully on the river bank.\",\n",
    "         \"the jewels were kept in a bank vault.\",\n",
    "         \"you can stow your jewellery away in the bank.\",\n",
    "         \"most of the money was in storage in bank vaults.\",\n",
    "         \"the diamonds are shut away in a bank vault somewhere.\",\n",
    "         \"thieves broke into the bank vault.\",\n",
    "         \"can I bank on your support?\",\n",
    "         \"you can bank on him to hand you a reasonable bill for your services.\",\n",
    "         \"don't bank on your friends to help you out of trouble.\",\n",
    "         \"you can bank on me when you need money.\",\n",
    "         \"i bank on your help.\"\n",
    "         ]\n",
    "from collections import OrderedDict\n",
    "context_embeddings = []\n",
    "context_tokens = []\n",
    "for sentence in sentences:\n",
    "  tokenized_text, tokens_tensor, segments_tensors = bert_text_preparation(sentence, tokenizer)\n",
    "  list_token_embeddings = get_bert_embeddings(tokens_tensor, segments_tensors, model)\n",
    "  # make ordered dictionary to keep track of the position of each   word\n",
    "  tokens = OrderedDict()\n",
    "  # loop over tokens in sensitive sentence\n",
    "  for token in tokenized_text[1:-1]:\n",
    "    # keep track of position of word and whether it occurs multiple times\n",
    "    if token in tokens:\n",
    "      tokens[token] += 1\n",
    "    else:\n",
    "      tokens[token] = 1\n",
    "  # compute the position of the current token\n",
    "    token_indices = [i for i, t in enumerate(tokenized_text) if t == token]\n",
    "    current_index = token_indices[tokens[token]-1]\n",
    "  # get the corresponding embedding\n",
    "    token_vec = list_token_embeddings[current_index]\n",
    "    \n",
    "    # save values\n",
    "    context_tokens.append(token)\n",
    "    context_embeddings.append(token_vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "filepath = os.path.join('inputs/')\n",
    "name = 'metadata_small.tsv'\n",
    "with open(os.path.join(filepath, name), 'w+') as file_metadata:\n",
    "  for i, token in enumerate(context_tokens):\n",
    "    file_metadata.write(token + '\\n')\n",
    "import csv\n",
    "name = 'embeddings_small.tsv'\n",
    "with open(os.path.join(filepath, name), 'w+') as tsvfile:\n",
    "    writer = csv.writer(tsvfile, delimiter='\\t')\n",
    "    for embedding in context_embeddings:\n",
    "        writer.writerow(embedding.numpy())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}