RでTransformerのAttention機構:Embedding

# パッケージの読み込み
library(reticulate)
library(ggplot2)

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

model_name <- "bert-base-uncased"
tokenizer <- transformers$AutoTokenizer$from_pretrained(model_name)
model <- transformers$AutoModel$from_pretrained(model_name)
# model$eval()

get_static_embedding <- function(word) {
  # 単語をトークンIDに変換
  token_id <- tokenizer$convert_tokens_to_ids(word)
  # 埋め込み行列から該当行を取り出す
  embedding_matrix <- model$embeddings$word_embeddings$weight$detach()$numpy()
  embedding_matrix[token_id + 1, ] # Python は0始まり、Rは1始まりなので+1
}

static_bank <- get_static_embedding("bank")

cat("=== 'bank' の静的 embedding (先頭6次元のみ表示) ===\n")
print(round(static_bank[1:6], 3))
cat(sprintf("次元数: %d\n", length(static_bank)))

get_contextual_embeddings <- function(sentence) {
  inputs <- tokenizer(sentence, return_tensors = "pt")

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

  # last_hidden_state: (1, seq_len, hidden_dim) -> (seq_len, hidden_dim)
  hidden <- outputs$last_hidden_state$squeeze()$numpy()

  tokens <- unlist(tokenizer$convert_ids_to_tokens(
    inputs$input_ids$squeeze()$tolist()
  ))

  list(tokens = tokens, embeddings = hidden)
}

# 指定した単語に対応するベクトルを取り出すヘルパー関数。
# BERT は WordPiece という方式でトークン化するため、
# 単語が複数のサブワードに分割される場合がある
# (例: "deposited" -> "deposit" + "##ed")。
# その場合は該当するサブワードのベクトルを平均する。

get_word_vector <- function(res, target_word) {
  target_word <- tolower(target_word)
  tokens <- res$tokens

  # ケース1: 単語がそのまま1トークンとして存在する
  exact_idx <- which(tokens == target_word)
  if (length(exact_idx) >= 1) {
    return(res$embeddings[exact_idx[1], ])
  }

  # ケース2: サブワードに分割されている場合、再構成しながら探す
  for (i in seq_along(tokens)) {
    if (!startsWith(tokens[i], "##") && startsWith(target_word, tokens[i])) {
      reconstructed <- tokens[i]
      idxs <- i
      j <- i
      while (reconstructed != target_word &&
        j < length(tokens) && startsWith(tokens[j + 1], "##")) {
        j <- j + 1
        reconstructed <- paste0(reconstructed, sub("^##", "", tokens[j]))
        idxs <- c(idxs, j)
      }
      if (reconstructed == target_word) {
        return(colMeans(res$embeddings[idxs, , drop = FALSE]))
      }
    }
  }

  stop(sprintf(
    "単語 '%s' が見つかりませんでした。トークン列: %s",
    target_word, paste(tokens, collapse = " ")
  ))
}

sentence_river <- "I sat on the river bank and watched the fish swim"
sentence_deposit <- "I put my money in the bank yesterday"
sentence_loan <- "The bank approved my loan application"

res_river <- get_contextual_embeddings(sentence_river)
res_deposit <- get_contextual_embeddings(sentence_deposit)
res_loan <- get_contextual_embeddings(sentence_loan)

bank_river <- get_word_vector(res_river, "bank") # 川岸の意味
bank_deposit <- get_word_vector(res_deposit, "bank") # 銀行の意味
bank_loan <- get_word_vector(res_loan, "bank") # 銀行の意味

cosine_sim <- function(a, b) sum(a * b) / (sqrt(sum(a^2)) * sqrt(sum(b^2)))

cat("\n=== 'bank' の文脈化 embedding 同士のコサイン類似度 ===\n")
cat(sprintf(
  "river-bank  vs  deposit-bank : %.3f  (意味が違う)\n",
  cosine_sim(bank_river, bank_deposit)
))
cat(sprintf(
  "deposit-bank vs  loan-bank    : %.3f  (意味が同じ)\n",
  cosine_sim(bank_deposit, bank_loan)
))
cat(sprintf(
  "river-bank  vs  loan-bank     : %.3f  (意味が違う)\n",
  cosine_sim(bank_river, bank_loan)
))

# 参考: 静的embeddingと文脈化embeddingの差も確認
cat(sprintf(
  "\n静的embedding vs river文脈でのbank   : %.3f\n",
  cosine_sim(static_bank, bank_river)
))
cat(sprintf(
  "静的embedding vs deposit文脈でのbank : %.3f\n",
  cosine_sim(static_bank, bank_deposit)
))

sentence_apple <- "I ate an apple for breakfast"
res_apple <- get_contextual_embeddings(sentence_apple)
word_apple <- get_word_vector(res_apple, "apple")

cat(sprintf("apple vs bank(river)   : %.3f\n", cosine_sim(word_apple, bank_river)))
cat(sprintf("apple vs bank(deposit) : %.3f\n", cosine_sim(word_apple, bank_deposit)))
cat(sprintf("apple vs bank(loan)    : %.3f\n", cosine_sim(word_apple, bank_loan)))