HuggingFace の BERT(AutoModel)の output から [CLS] 埋め込みを得る

この部分毎回忘れてググっている気がするのでメモ。

outputs["last_hidden_state"][:, 0, :]

でOK（入力文数×768 の Tensor になるはず）。

入力に一文しか入れていなくて直接ベクトルを取り出したい場合は

output["last_hidden_state"][0, 0, :]

---

追記:

outputs[1]

で得られる埋め込みは、pooling 層を通した後の埋め込みなので、[CLS] に対応する最終隠れ層を得たい場合はこちらを使うべきではない