Cross-encoder model (from HF): "cross-encoder/ms-marco-MiniLM-L-6-v2"

[ndrean]

elixir-nx/bumblebee#251

Bumblebee did not yet implement the architecture :for_sequence_classification.

efmodule RAG.CrossEncoder do
  def load_model do
    repo= {:hf, "cross-encoder/ms-marco-MiniLM-L-6-v2"}
    tokenizer = {:hf, "bert-base-uncased"}
    {:ok, model_info} = Bumblebee.load_model(repo)
    {:ok, tokenizer} = Bumblebee.load_tokenizer(tokenizer)

    {model_info, tokenizer}
  end

  def rerank(query, documents) do
    # Prepare input pairs for cross-encoder
    {model_info, tokenizer} = load_model()
    input_pairs = 
     Bumblebee.apply_tokenizer(tokenizer, 
       Enum.map(documents, fn doc -> 
         {query, doc.content}
       end)
     )

    # Run cross-encoder in batches
    outputs = Axon.predict(model_info.model, model_info.params, input_pairs)
     

    # Combine scores with original documents and sort
    Enum.zip(documents, outputs.logits |> Nx.to_flat_list())
    |> Enum.sort_by(fn {_, score} -> score end, :desc)
    |> Enum.map(fn {doc, _} -> doc.content end)
  end
end

# Load the model
cross_encoder_serving = RAG.CrossEncoder.load_model()


# Rerank the results obtained from the HSWN (first creaming step)
reranked_results = RAG.CrossEncoder.rerank(query, top_results, cross_encoder_serving)