検索のシンプルな実装を作ってみる系の話です。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とかにインデックスして、検索して試せるようにするとかはできそうな気がするので、これはやってみたい気はしてます。