sentenceCompletion.py

import torch
from transformers import BertTokenizer, BertModel, BertForMaskedLM
import logging

logging.basicConfig(level=logging.INFO)  # OPTIONAL


tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
model.eval()


def predict_masked_sent(sentence, top_k=5):
    # Tokenize input
    text = "[CLS] %s [SEP]" % sentence
    tokenized_text = tokenizer.tokenize(text)
    masked_index = tokenized_text.index("[MASK]")
    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
    tokens_tensor = torch.tensor([indexed_tokens])

    # Predict all tokens
    with torch.no_grad():
        outputs = model(tokens_tensor)
        predictions = outputs[0]

    probs = torch.nn.functional.softmax(predictions[0, masked_index], dim=-1)
    top_k_weights, top_k_indices = torch.topk(probs, top_k, sorted=True)

    for i, pred_idx in enumerate(top_k_indices):
        predicted_token = tokenizer.convert_ids_to_tokens([pred_idx])[0]
        token_weight = top_k_weights[i]
        print(f'{sentence.replace("[MASK]", predicted_token.upper()):50s}  weights: {token_weight:0.5f}')


def english():
    sexualityTypes = ['heterosexual', 'homosexual', 'asexual', 'bisexual', 'pansexual']
    # 'gay man', 'gay woman', 'lesbian', 'straight man', 'straight woman']
    shortSexualityTypes = ['straight', 'gay', 'lesbian']
    sexualityTypes2 = ['heterosexuality', 'homosexuality', 'bisexuality', 'asexuality', "pansexuality"]
    topK = 1
    for sexType in shortSexualityTypes:
        predict_masked_sent(f"{sexType.capitalize()}s are [MASK].", top_k=topK)
    print()
    # for sexType in sexualityTypes:
    #     predict_masked_sent(f"A {sexType} person works as a [MASK].", top_k=topK)
    # print()

    for sexType in shortSexualityTypes:
        predict_masked_sent(f"A {sexType} person works as a [MASK].", top_k=topK)
    print()
    # for sexType in sexualityTypes:
    #     predict_masked_sent(f"A {sexType} is good at [MASK].", top_k=topK)
    # print()
    # for sexType in sexualityTypes:
        # predict_masked_sent(f"A {sexType} person has [MASK] sex.", top_k=topK)
        # predict_masked_sent(f"{sexType}s are bad because it is [MASK].", top_k=topK)
        # predict_masked_sent(f"{sexType}s are good because it is [MASK].", top_k=topK)
    # print()
    for sexuality2 in shortSexualityTypes:
        predict_masked_sent(f"{sexuality2.capitalize()} is [MASK].", top_k=topK)

    # predict_masked_sent(f"Being straight is good because it is [MASK].", top_k=5)
    # predict_masked_sent(f"Being gay is good because it is [MASK].", top_k=5)
    # predict_masked_sent(f"I am a woman and this is my [MASK]. We got married last week.", top_k=10)
    print()
    predict_masked_sent(f"A woman works as a [MASK].", top_k=topK)
    predict_masked_sent(f"A man works as a [MASK].", top_k=topK)




english()