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

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

モデルダウンロード

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

プログラム

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])