# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="kY1SDwkhyRFN" colab_type="text"
# https://colab.research.google.com/drive/11ScUROhQM6DznUe4qqJB3klfuXywjmTQ?usp=sharing

# + id="vDftFqnIt4WH" colab_type="code" colab={}
# ! pip install -q torch
# ! pip install -q pytorch-pretrained-bert

# + id="w1NJeyxst5Eq" colab_type="code" colab={}
import torch
from pytorch_pretrained_bert import BertTokenizer, BertModel, BertForMaskedLM

# Load pre-trained model tokenizer (vocabulary)
modelpath = "bert-base-uncased"
tokenizer = BertTokenizer.from_pretrained(modelpath)

# + id="1MgGX1oJuJ-X" colab_type="code" outputId="4643923f-8b23-4837-d4b0-84e81538cfaf" colab={"base_uri": "https://localhost:8080/", "height": 306}

text = "dummy. although he had already eaten a large meal, he was still very hungry."
target = "hungry"
tokenized_text = tokenizer.tokenize(text)
tokenized_text

# + id="YPn7noY_uMik" colab_type="code" outputId="2be52148-4c16-4051-b257-9b0d5cdb5bdd" colab={"base_uri": "https://localhost:8080/", "height": 306}

# Mask a token that we will try to predict back with `BertForMaskedLM`
masked_index = tokenized_text.index(target)
tokenized_text[masked_index] = '[MASK]'
tokenized_text

# + id="xpMqHToYuOqM" colab_type="code" outputId="8dcfba06-2d17-4b82-bdd6-cbef147a6b2e" colab={"base_uri": "https://localhost:8080/", "height": 306}
# Convert token to vocabulary indices
indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
indexed_tokens

# + id="Vv6wY7VBuQba" colab_type="code" outputId="491252e9-125d-4cef-a667-4b1de36dc215" colab={"base_uri": "https://localhost:8080/", "height": 34}
# Define sentence A and B indices associated to 1st and 2nd sentences (see paper)
segments_ids = [1] * len(tokenized_text)
segments_ids

# + id="UdfJ1R9GuSJG" colab_type="code" outputId="4d407f94-6ab5-4392-8a61-3ae8038a358b" colab={"base_uri": "https://localhost:8080/", "height": 34}
# this is for the dummy first sentence. 
segments_ids[0] = 0
segments_ids[1] = 0
segments_ids

# + id="_-t-aL1uuT0G" colab_type="code" outputId="9b989be5-37a8-40f6-fa82-c6ebca8e213e" colab={"base_uri": "https://localhost:8080/", "height": 68}

# Convert inputs to PyTorch tensors
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])
tokens_tensor, segments_tensors

# + id="ay7DOOv5uVRD" colab_type="code" colab={}
# Load pre-trained model (weights)
model = BertForMaskedLM.from_pretrained(modelpath)
model.eval()


# + id="clzecg6guWpM" colab_type="code" outputId="b95a987e-7088-429e-d1f1-f44239edb7dc" colab={"base_uri": "https://localhost:8080/", "height": 34}

# Predict all tokens
predictions = model(tokens_tensor, segments_tensors)
predicted_index = torch.argmax(predictions[0, masked_index]).item()
predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])
predicted_index, predicted_token

# + id="-yOHWyt-uYzj" colab_type="code" outputId="734fc21e-19ad-4717-c401-562de79ecc97" colab={"base_uri": "https://localhost:8080/", "height": 85}

print("Original:", text)
print("Masked:", " ".join(tokenized_text))

print("Predicted token:", predicted_token)
print("Other options:")


# + id="x3OghyKGuacA" colab_type="code" outputId="273c5a0c-7af2-4e85-bd6b-a5fa79708b5e" colab={"base_uri": "https://localhost:8080/", "height": 187}
# just curious about what the next few options look like.
for i in range(10):
    predictions[0,masked_index,predicted_index] = -11100000
    predicted_index = torch.argmax(predictions[0, masked_index]).item()
    predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])
    print(predicted_token)

StarcoderdataJupyterScriptsDedupFiltered