「pfnet/plamo-embedding-1b」をChromaDBのEmbeddingで使ってみた

Python Webアプリ 自然言語処理

お疲れ様です。

今日の昼頃Preferred Networks（PFN）からPLaMoをベースにしたテキスト埋め込みモデル（Embeddingモデル）が公開されていました！

tech.preferred.jp

この公開された「plamo-embedding-1b」ですが、他の日本語テキスト埋め込みモデルと比較してもかなり性能が良いとのこと。
このモデル自体はHugging Faceで公開されており、ライセンスもApache v2.0ライセンスなので商用利用可。めちゃくちゃありがたいです。

huggingface.co

とにかくまずは使ってみたいということで、ChromaDBのEmbeddingモデルとして使ってみました。
作成したデータベースを使ってRAGを動かすところまでやってみます。 ソースコードについてはStreamlitで作成しているチャットボットアプリの一部になります。

github.com

基本的にはこれまでのEmbeddingモデル設定をplamo-embedding-1bに置き換えるだけです。
これまではintfloat/multilingual-e5-baseを使用していましたがここを置き換えました。さすがにモデルサイズは大きくなっておりCPUだけの処理では時間がかかりすぎたのでCUDAを使用しています。
データベース作成の部分とRAG実行時の入力文のベクトル化の部分をこちらに置き換えています。

# ベクトル化する準備
model_kwargs = {"device": "cuda", "trust_remote_code": True}
embedding = HuggingFaceEmbeddings(
    model_name="pfnet/plamo-embedding-1b",
    model_kwargs=model_kwargs
)

アプリ側でRAGを実行してみました。データベースは急造で大した大きさではないですので参考程度に。
データベースにしたのはエレファントカシマシに関するWikipediaのテキスト情報です。 おそらく以下のWikipediaの内容が取得されていますが、うまく指定した曲の内容だけを取れているように思いますね。 ちょっとまだサンプルが少ないのでもっと試していきたいところです！
悲しみの果て - Wikipedia