RでTransformerのAttention機構:BERT

library(reticulate)

transformers <- import("transformers")
torch <- import("torch")

model_name <- "bert-base-uncased"
tokenizer <- transformers$AutoTokenizer$from_pretrained(model_name)

# MLM(穴埋め)タスク用のヘッド付きモデルを読み込む。
# 内部的には encoder 部分 (model$bert) と、
# 語彙全体への予測ヘッド (model$cls) を両方持っている
model <- transformers$BertForMaskedLM$from_pretrained(model_name)
# model$eval()

config <- model$config

# 総パラメータ数、および埋め込み部分とEncoder部分の内訳を確認
count_params <- function(module) {
  total <- 0
  for (p in iterate(module$parameters())) {
    total <- total + as.integer(p$numel())
  }
  total
}

embedding_params <- count_params(model$bert$embeddings)
encoder_params <- count_params(model$bert$encoder)
total_params <- count_params(model)

bert_summary <- data.frame(
  項目 = c(
    "語彙数 (vocab_size)",
    "埋め込み次元 (hidden_size)",
    "Encoder層数 (num_hidden_layers)",
    "Attention head数 (1層あたり)",
    "1 head あたりの次元",
    "Feed-forward層の中間次元",
    "最大入力トークン数",
    "埋め込み層のパラメータ数",
    "Encoder(12層)のパラメータ数",
    "モデル全体のパラメータ数"
  ),
  値 = c(
    format(config$vocab_size, big.mark = ","),
    format(config$hidden_size, big.mark = ","),
    format(config$num_hidden_layers, big.mark = ","),
    format(config$num_attention_heads, big.mark = ","),
    format(as.integer(config$hidden_size / config$num_attention_heads), big.mark = ","),
    format(config$intermediate_size, big.mark = ","),
    format(config$max_position_embeddings, big.mark = ","),
    format(embedding_params, big.mark = ","),
    format(encoder_params, big.mark = ","),
    format(total_params, big.mark = ",")
  )
)

knitr::kable(bert_summary,
  format = "html", align = c("l", "r"),
  caption = "BERT-base のアーキテクチャ情報"
)

softmax <- function(x) {
  x_shifted <- x - max(x)
  exp(x_shifted) / sum(exp(x_shifted))
}

get_top_predictions <- function(sentence, tokenizer, model, top_k = 5) {
  inputs <- tokenizer(sentence, return_tensors = "pt")

  with(torch$no_grad(), {
    outputs <- model(
      input_ids = inputs$input_ids,
      attention_mask = inputs$attention_mask
    )
  })

  logits <- outputs$logits$squeeze()$detach()$numpy() # (seq_len, vocab_size)

  token_ids <- inputs$input_ids$squeeze()$tolist()
  mask_pos <- which(unlist(token_ids) == tokenizer$mask_token_id)

  mask_logits <- logits[mask_pos, ]
  probs <- softmax(mask_logits)

  top_idx <- order(probs, decreasing = TRUE)[1:top_k]
  top_token_ids <- as.list(as.integer(top_idx - 1)) # Rの1-indexを Pythonの0-indexに戻す
  top_tokens <- unlist(tokenizer$convert_ids_to_tokens(top_token_ids))

  data.frame(token = top_tokens, probability = probs[top_idx])
}

sentence1 <- sprintf("The capital of France is %s.", tokenizer$mask_token)
cat(sprintf("=== MLM予測: \"%s\" ===\n", sentence1))
pred1 <- get_top_predictions(sentence1, tokenizer, model)
print(pred1)

sentence2 <- sprintf("I went to the %s to withdraw some cash.", tokenizer$mask_token)
cat(sprintf("=== 双方向性の確認: \"%s\" ===\n", sentence2))

inputs2 <- tokenizer(sentence2, return_tensors = "pt")
token_ids2 <- unlist(inputs2$input_ids$squeeze()$tolist())
mask_pos2 <- which(token_ids2 == tokenizer$mask_token_id)
seq_len2 <- length(token_ids2)

# --- 通常 (双方向): attention_mask は全て1のまま ---
attn_full <- inputs2$attention_mask

# --- 右側を隠す (擬似的な単方向化): [MASK]より後ろの
#     トークンへのattentionを0にする ---
attn_restricted_vec <- rep(1L, seq_len2)
if (mask_pos2 < seq_len2) {
  attn_restricted_vec[(mask_pos2 + 1):seq_len2] <- 0L
}
attn_restricted <- torch$tensor(list(as.integer(attn_restricted_vec)), dtype = torch$long)

run_mlm <- function(input_ids, attention_mask, mask_pos) {
  with(torch$no_grad(), {
    outputs <- model(input_ids = input_ids, attention_mask = attention_mask)
  })
  logits <- outputs$logits$squeeze()$detach()$numpy()
  softmax(logits[mask_pos, ])
}

probs_full <- run_mlm(inputs2$input_ids, attn_full, mask_pos2)
probs_restricted <- run_mlm(inputs2$input_ids, attn_restricted, mask_pos2)

bank_id <- tokenizer$convert_tokens_to_ids("bank")

cat("\n--- 通常(双方向、右側の文脈あり) ---\n")
top_full <- order(probs_full, decreasing = TRUE)[1:5]
print(data.frame(
  token = unlist(tokenizer$convert_ids_to_tokens(as.list(as.integer(top_full - 1)))),
  probability = probs_full[top_full]
))
cat(sprintf("'bank' に割り当てられた確率: %.4f\n", probs_full[bank_id + 1]))

cat("\n--- 右側の文脈を隠した場合(擬似的な単方向) ---\n")
top_restricted <- order(probs_restricted, decreasing = TRUE)[1:5]
print(data.frame(
  token = unlist(tokenizer$convert_ids_to_tokens(as.list(as.integer(top_restricted - 1)))),
  probability = probs_restricted[top_restricted]
))
cat(sprintf("'bank' に割り当てられた確率: %.4f\n", probs_restricted[bank_id + 1]))