sbintuitions / JMTEB

The evaluation scripts of JMTEB (Japanese Massive Text Embedding Benchmark)

Creative Commons Attribution Share Alike 4.0 International

24 stars 4 forks source link

[Feature] Support customizing `model_kwargs` #41

Closed lsz05 closed 2 months ago

lsz05 commented 2 months ago

関連する Issue / PR

39

PR をマージした後の挙動の変化

TextEmbedderに出力型flag (sentence transformersに倣ってconvert_to_tensor, convert_to_numpy)をメンバーに入れる
モデルを初期化する際にmodel_kwargsを設定できるようになった
jsonargparse対応のため，model_kwargs引数を読み込む関数を作成
minor fixes
テストケースを追加

挙動の変更を達成するために行ったこと

TextEmbedderに出力型flag (sentence transformersに倣ってconvert_to_tensor, convert_to_numpy)をメンバーに入れる
モデルを初期化する際にmodel_kwargs設定（特に推論時のdtype，そのための読み込み関数を作成）を読み込む
後処理にtorch.Tensor型使用のシナリオ（検索とreranking），Tensor→ndarray→Tensorの変換をやめて，Tensorそのまま出力させる
torch.float16を指定する際の動作確認をテストケースとして追加する

動作確認

[x] テストが通ることを確認した
[x] マージ先がdevブランチであることを確認した