Zero-shot classification

Mix.install([
  {:req, "~> 0.5"},
  {:hnswlib, "~> 0.1"},
  {:instructor, "~> 0.1"},
  {:bumblebee, "~> 0.6.0"},
  {:exla, "~> 0.9"},

  {:elixir_make, "~> 0.9", override: true}
], 
  config: [nx: [default_backend: EXLA.Backend]]
)

Application.put_env(:exla, :clients, cuda: [platform: :cuda, preallocate: false])

Setup

client = Req.new(base_url: "https://models.arrakis.upmaru.network")

defmodule QueryClassification do
  use Ecto.Schema
  use Instructor

  @llm_doc """
  Query classification
  
  ## Fields:
  - reply: Should the reply be list of results or a recommendation?
  - about: What is the query about? Location, character or general knowledge?
  """
  @primary_key false
  embedded_schema do
    field :reply, Ecto.Enum, values: [:results, :recommendation]
    field :about, Ecto.Enum, values: [:general, :character, :location]
  end
end

With embeddings

space = :cosine
dimensions = 1024
max_elements = 3

{:ok, index} = HNSWLib.Index.new(space, dimensions, max_elements)

body = %{
  input: ["query with no specific category", "query about movie character", "query related to the location or the setting of the movie"],
  model: "multilingual-e5-large",
  encoding_format: "float"
}
  
{:ok, %{body: response}} = Req.post(client, url: "/v1/embeddings", json: body)

%{"data" => embeddings} = response

embeddings = 
  Enum.sort_by(embeddings, fn e -> 
    e["index"]
  end)
  |> Enum.map(fn e -> e["embedding"] end)
  |> Nx.tensor()

HNSWLib.Index.add_items(index, embeddings)

original_query = "query: movies involving robots"

query_body = %{
  input: original_query,
  model: "multilingual-e5-large",
  encoding_format: "float"
}

{:ok, %{body: query}} = Req.post(client, url: "/v1/embeddings", json: query_body)

%{"data" => [%{"embedding" => embedding}]} = query

query = Nx.tensor(embedding)

{:ok, labels, dists} = HNSWLib.Index.knn_query(index, query, k: 3)

With reranking model

original_query = "query: movies involving robots"

body = %{
  query: original_query,
  documents: ["query is general", "query about characters in the movie", "query related to the location or the setting of the movie"],
  model: "bge-reranker-v2-m3"
}

{:ok, %{body: %{"results" => reranking_results}}} = Req.post(client, url: "/v1/rerank", json: body)

With LLM

instructor_config = [
  adapter: Instructor.Adapters.Llamacpp,
  api_url: "https://models.arrakis.upmaru.network"
]

{:ok, setting} = Instructor.chat_completion([
    model: "mistral-small-24b",
    mode: :json_schema,
    response_model: QueryClassification,
    messages: [
      %{
        role: "user",
        content: """
           Correctly identify what the query is about and what type of reply should be made
        
           
           Movies that is set in the sea or the ocean.
           
          """
      }
    ]
  ], instructor_config)

Using Zero-shot classification model

{:ok, model} = Bumblebee.load_model({:hf, "facebook/bart-large-mnli"})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "facebook/bart-large-mnli"})

categories = ["general", "character", "location"]
categories_serving = Bumblebee.Text.zero_shot_classification(model, tokenizer, categories)

output = Nx.Serving.run(categories_serving, "Movies that take place in the mountains")