カテゴリー: Machine Learning

検索のシンプルな実装を作ってみる系の話です。simple-spladeとか、ちょこちょこ作れるときに作ってみたりしているけど、基本に返って、シンプルなDenseなベクトル検索用の実装を作ってみるかな、と思い始めたのがsimple-dense-retrievalです。

いつものように利用しているデータはESCIのAmazonの商品検索データです。テキストのembeddingには、multilingual-e5-smallをベースにモデルを作る感じです。

ESCIのデータをダウンロードしておけば、python main.pyで学習してモデルができる感じです。query_encoderとdocuemnt_encoderという形で、モデルを保存します。

推論用APIも用意しておこうと思い、今回はTorchServeを利用してます。torchserveの中に置いてあります。query_encoderとdocument_encoderのMARファイルを作って、それをTorchServeのDockerで読み込んで利用するみたいな感じです。

モデルのチューニングの余地は結構ありそうな気がしますが、気が向いたときに調整するかもしれませんが、simple-〜は仕組みの理解のための作っている感じなので、これをベースにいろいろと考えるための叩き台みたいな感じの想定です。

そんな感じですが、simple-dense-retrievalで他にやってみたいこととしては、TorchServe以外にも推論APIはあるので、Triton Inference Serverとかでも動くようにしたいところではあります。ですが、これをするには、model.pyをリファクタリングしないと無理そうなので、これは諦め気味…。他には、実際にOpenSearchとかにインデックスして、検索して試せるようにするとかはできそうな気がするので、これはやってみたい気はしてます。

simple-spladeを作った勢いで、ChatGPTと相談しながら、simple-*的な派生品を作り始めているのですが、ひとまず、simple-spladeppとして、SPLADE++を作ってみた。

teacherモデルをどうするかはあるけど、ざっとみた感じでは、こんな感じかなと思われる。日本語をどうするか？を考えるなら、ベースモデルを作るところも考えた方が良いような気もするけど、それは気が向いたら考えることにする。

そんな感じで、今後もsimple-*シリーズを増やしてみようかなと。

naverのSPLADEの実装は、GitHubのnaver/spladeにあるけど、これが商用不可だったりするので、ゼロから実装して、何かで使いたいときの際にベースとして利用できるものがほしいなーと思ってました。というわけで、シンプルなSPLADEの実装を作ろうということで、marevol/simple-spladeとして、作ってみました。ChatGPTと壁打ちしながら、作っていった感じではあるけど、たぶん、それっぽい感じにはなったかなと思います。

main.pyでAmazonのECSIデータを利用して、クエリー(query)と商品名(product_title)を学習する感じになってます。ExactのEをpositiveで、それ以外をnegativeとして扱ってます。コード量も多くないので、簡単に眺めることができるのではないかな…と思いますし、たぶん、手を加えやすいと思います。

マルチリンガルで、とChatGPTに言ったら、xlm-roberta-baseになったので、そのままですが、別なモデルでも良いような気もしてます…。

勢いで作った感じはありますが、何かあれば、修正してください。