PLaMo-Embedding-1Bを触ってみる文章を数値ベクトルに変換することができるテキスト埋め込みモデルのPLaMo-Embedding-1Bを今回触ってみます

文章を数値ベクトルに変換することができるテキスト埋め込みモデルのPLaMo-Embedding-1Bを今回触ってみます

pfnet/plamo-embedding-1b · Hugging Face

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

前提条件
モデルダウンロード
プログラム
実行

前提条件

Python
git
git lfs(https://git-lfs.com/)

が使用するPCにインストールされていることを確認してください

モデルダウンロード

とりあえずgitコマンドでPLaMo-Embedding-1Bをcloneしてきます

以下のコマンドをPowershellかコマンドプロンプトに貼り付けて実行します

git clone https://huggingface.co/pfnet/plamo-embedding-1b

プログラム

PLaMo-Embedding-1Bのフォルダの上位にmain.pyを作成します

動作確認のためサンプルコードをmain.pyにいれます

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

# You can download models from the Hugging Face Hub 🤗 as follows:
tokenizer = AutoTokenizer.from_pretrained("pfnet/plamo-embedding-1b", trust_remote_code=True)
model = AutoModel.from_pretrained("pfnet/plamo-embedding-1b", trust_remote_code=True)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

query = "PLaMo-Embedding-1Bとは何ですか？"
documents = [
    "PLaMo-Embedding-1Bは、Preferred Networks, Inc. によって開発された日本語テキスト埋め込みモデルです。",
    "最近は随分と暖かくなりましたね。"
]

with torch.inference_mode():
    # For embedding query texts in information retrieval, please use the `encode_query` method.
    # You also need to pass the `tokenizer`.
    query_embedding = model.encode_query(query, tokenizer)
    # For other texts/sentences, please use the `encode_document` method.
    # Also, for applications other than information retrieval, please use the `encode_document` method.
    document_embeddings = model.encode_document(documents, tokenizer)

# The similarity between vectors obtained by inputting sentences into the model is high for similar sentences and low for dissimilar sentences.
# This feature can be utilized for applications such as information retrieval.
similarities = F.cosine_similarity(query_embedding, document_embeddings)
print(similarities)
# tensor([0.8812, 0.5533])

Pythonで必要なパッケージをインストールします

pip install torch torchvision transformers

以下の通りに実行すると動作しました

実行

python .\main.py
model.safetensors: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2.10G/2.10G [09:23<00:00, 3.73MB/s]
tensor([0.8812, 0.5533])

以下は実行結果です。これは、クエリの埋め込みベクトルと各文書の埋め込みベクトルとの間のコサイン類似度を計算します。