Merge pull request #47 from ndrean/semantic

Semantic search added

Author: LuchoTurtle

.gitignore

@@ -19,6 +19,9 @@ erl_crash.dump
 # Also ignore archive artifacts (built via "mix archive.build").
 *.ez
 
+# Ignore DB dumps
+*.db
+
 # Temporary files, for example, from tests.
 /tmp/
 
@@ -37,4 +40,8 @@ npm-debug.log
 
 # Bumblebee model directory
 .bumblebee/*
-.elixir_ls
+.elixir_ls
+
+# KNN index direcotry
+priv/static/uploads/indexes.bin
+

README.md

@@ -12,13 +12,16 @@ defmodule Comparison.Models do
   def verify_and_download_model(model, force_download? \\ false) do
     case force_download? do
       true ->
-        File.rm_rf!(model.cache_path) # Delete any cached pre-existing model
-        download_model(model)         # Download model
+        # Delete any cached pre-existing model
+        File.rm_rf!(model.cache_path)
+        # Download model
+        download_model(model)
 
       false ->
         # Check if the model cache directory exists or if it's not empty.
         # If so, we download the model.
         model_location = Path.join(model.cache_path, "huggingface")
+
         if not File.exists?(model_location) or File.ls!(model_location) == [] do
           download_model(model)
         end
@@ -50,7 +53,7 @@ defmodule Comparison.Models do
   # It will load the model and the respective the featurizer, tokenizer and generation config if needed,
   # and return a map with all of these at the end.
   defp load_offline_model_params(model) do
-    Logger.info("Loading #{model.name}...")
+    Logger.info("ℹ️ Loading #{model.name}...")
 
     # Loading model
     loading_settings = {:hf, model.name, cache_dir: model.cache_path, offline: true}
@@ -92,7 +95,7 @@ defmodule Comparison.Models do
   # Downloads the models according to a given %ModelInfo struct.
   # It will load the model and the respective the featurizer, tokenizer and generation config if needed.
   defp download_model(model) do
-    Logger.info("Downloading #{model.name}...")
+    Logger.info("ℹ️ Downloading #{model.name}...")
 
     # Download model
     downloading_settings = {:hf, model.name, cache_dir: model.cache_path}

_comparison/manage_models.exs

@@ -80,7 +80,7 @@ defmodule Benchmark do
     coco_dataset_images_path = File.cwd!() |> Path.join("coco_dataset") |> Path.join("*.jpg")
     files = Path.wildcard(coco_dataset_images_path)
 
-    #coco_dataset_captions =
+    # coco_dataset_captions =
     #  File.stream!(File.cwd!() |> Path.join("coco_dataset") |> Path.join("captions.csv"))
     #  |> CSV.decode!()
     #  |> Enum.map(& &1)
@@ -120,7 +120,7 @@ defmodule Benchmark do
 
     # Go over each image and make prediction
     Enum.each(vips_images_with_captions, fn image ->
-      Logger.info("Benchmarking image #{image.id}...")
+      Logger.info("📊 Benchmarking image #{image.id}...")
 
       # Run the prediction
       {time_in_microseconds, prediction} =

_comparison/run.exs

@@ -11,23 +11,21 @@ export default {
       blue = ["bg-blue-500", "hover:bg-blue-700"],
       pulseGreen = ["bg-green-500", "hover:bg-green-700", "animate-pulse"];
 
-
     _this = this;
 
     // Adding event listener for "click" event
     recordButton.addEventListener("click", () => {
-
       // Check if it's recording.
       // If it is, we stop the record and update the elements.
       if (mediaRecorder && mediaRecorder.state === "recording") {
         mediaRecorder.stop();
+        // audioChunks.getAudioTracks()[0].stop();
         text.textContent = "Record";
-      } 
+      }
 
       // Otherwise, it means the user wants to start recording.
       else {
         navigator.mediaDevices.getUserMedia({ audio: true }).then((stream) => {
-
           // Instantiate MediaRecorder
           mediaRecorder = new MediaRecorder(stream);
           mediaRecorder.start();
@@ -39,7 +37,7 @@ export default {
 
           // Add "dataavailable" event handler
           mediaRecorder.addEventListener("dataavailable", (event) => {
-            audioChunks.push(event.data);
+            event.data.size > 0 && audioChunks.push(event.data);
           });
 
           // Add "stop" event handler for when the recording stops.
@@ -57,4 +55,4 @@ export default {
       }
     });
   },
-};
+};

assets/js/micro.js

@@ -13,7 +13,7 @@ config :app,
   generators: [timestamp_type: :utc_datetime]
 
 # Tells `NX` to use `EXLA` as backend
-# config :nx, default_backend: EXLA.Backend 
+# config :nx, default_backend: EXLA.Backend
 # needed to run on `Fly.io`
 config :nx, :default_backend, {EXLA.Backend, client: :host}
 

config/config.exs

@@ -10,7 +10,6 @@ config :app, App.Repo,
   show_sensitive_data_on_connection_error: true,
   pool_size: 10
 
-
 # For development, we disable any cache and enable
 # debugging and code reloading.
 #

config/dev.exs

@@ -26,7 +26,8 @@ config :logger, level: :warning
 # Initialize plugs at runtime for faster test compilation
 config :phoenix, :plug_init_mode, :runtime
 
-
 # App configuration
 config :app,
+  start_genserver: false,
+  knnindex_indices_test: true,
   use_test_models: true

config/test.exs

@@ -926,7 +926,7 @@ defmodule App.Models do
   # It will load the model and the respective the featurizer, tokenizer and generation config if needed,
   # and return a map with all of these at the end.
   defp load_offline_model(model) do
-    Logger.info("Loading #{model.name}...")
+    Logger.info("ℹ️ Loading #{model.name}...")
 
     # Loading model
     loading_settings = {:hf, model.name, cache_dir: model.cache_path, offline: true}
@@ -968,7 +968,7 @@ defmodule App.Models do
   # Downloads the models according to a given %ModelInfo struct.
   # It will load the model and the respective the featurizer, tokenizer and generation config if needed.
   defp download_model(model) do
-    Logger.info("Downloading #{model.name}...")
+    Logger.info("ℹ️ Downloading #{model.name}...")
 
     # Download model
     downloading_settings = {:hf, model.name, cache_dir: model.cache_path}

deployment.md

@@ -0,0 +1,127 @@
+Mix.install([
+{:bumblebee, "~> 0.4.2"},
+{:exla, "~> 0.6.4"},
+{:nx, "~> 0.6.4 "},
+{:hnswlib, "~> 0.1.4"}
+])
+
+Nx.global_default_backend(EXLA.Backend)
+
+{:ok, index} = HNSWLib.Index.new(_space = :cosine, _dim = 384, _max_elements = 200)
+transformer = "sentence-transformers/paraphrase-MiniLM-L6-v2"
+{:ok, %{model: _model, params: _params} = model_info} =
+      Bumblebee.load_model({:hf, transformer})
+
+{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, transformer})
+serving = Bumblebee.Text.TextEmbedding.text_embedding(
+      model_info, 
+      tokenizer, 
+      defn_options: [compiler: EXLA, lazy_transfers: :never]
+      #output_pool: :mean_pooling,
+      #output_attribute: :hidden_state,
+      #embedding_processor: :l2_norm,
+    )
+
+%{embedding: data} = Nx.Serving.run(serving, "small") |>dbg()
+HNSWLib.Index.add_items(index, data)
+HNSWLib.Index.get_count(index) |> dbg()
+
+%{embedding: data} = Nx.Serving.run(serving, "tall") |> dbg()
+HNSWLib.Index.add_items(index, data)
+HNSWLIb.Index.get_count(index) |> dbg()
+
+%{embedding: data} = Nx.Serving.run(serving, "high")
+{:ok, labels, distances} = HNSWLib.Index.knn_query(index, data, k: 1) |> dbg()
+idx = Nx.to_flat_list(labels[0])
+{:ok, dt} = HNSWLib.Index.get_items(index, idx)
+Nx.stack(Enum.map(dt, fn d -> Nx.from_binary(d, :f32) end))
+
+defmodule Embedding do
+  use GenServer
+  @indexes "indexes.bin"
+
+  def start_link(norm) do
+    GenServer.start_link(__MODULE__, norm, name: __MODULE__)
+  end
+
+  # upload or create a new index file
+  def init(norm) do
+    space = norm
+
+    {:ok, index} =
+      case File.exists?(@indexes) do
+        false ->
+          HNSWLib.Index.new(_space = space, _dim = 384, _max_elements = 200)
+
+        true ->
+          HNSWLib.Index.load_index(space, 384, @indexes)
+      end
+
+    model_info = nil
+    tokenizer = nil
+    {:ok, {model_info, tokenizer, index}, {:continue, :load}}
+  end
+
+  def handle_continue(:load, {_, _, index}) do
+    transformer = "sentence-transformers/paraphrase-MiniLM-L6-v2"
+
+    {:ok, %{model: _model, params: _params} = model_info} =
+      Bumblebee.load_model({:hf, transformer})
+
+    {:ok, tokenizer} =
+      Bumblebee.load_tokenizer({:hf, transformer})
+
+    {:noreply, {model_info, tokenizer, index}}
+  end
+
+  def serve() do
+    GenServer.call(__MODULE__, :serve)
+  end
+
+  def get_count do
+    GenServer.call(__MODULE__, :get_count)
+  end
+
+  def get_index do
+    GenServer.call(__MODULE__, :get_index)
+  end
+
+  def handle_call(:serve, _from, {model_info, tokenizer, index} = state) do
+    serving = Bumblebee.Text.TextEmbedding.text_embedding(
+      model_info, 
+      tokenizer, 
+      output_pool: :mean_pooling,
+      output_attribute: :hidden_state,
+      embedding_processor: :l2_norm,
+      defn_options: [compiler: EXLA, lazy_transfers: :never]
+    )
+    {:reply, {serving, index}, state}
+  end
+
+  def handle_call(:get_count, _, {_, _, index} = state) do
+    {:ok, count} = HNSWLib.Index.get_current_count(index)
+    {:reply, count, state}
+  end
+
+  def handle_call(:get_index, _, {_, _, index} = state) do
+    {:reply, index, state}
+  end
+end
+
+{:ok, pid} = GenServer.start_link(Embedding, :l2)
+
+{serving, index} = GenServer.call(pid, :serve)
+
+%{embedding: data} = Nx.Serving.run(serving, "small") |>dbg()
+HNSWLib.Index.add_items(index, data)
+GenServer.call(pid, :get_count) |> dbg()
+
+%{embedding: data} = Nx.Serving.run(serving, "tall") |> dbg()
+HNSWLib.Index.add_items(index, data)
+GenServer.call(pid, :get_count) |> dbg()
+
+%{embedding: data3} = Nx.Serving.run(serving, "high")
+{:ok, labels, distances} = HNSWLib.Index.knn_query(index, data, k: 1) |> dbg()
+idx = Nx.to_flat_list(labels[0])
+{:ok, dt} = HNSWLib.Index.get_items(index, idx)
+Nx.stack(Enum.map(dt, fn d -> Nx.from_binary(d, :f32) end))

example.txt

@@ -5,10 +5,28 @@ defmodule App.Application do
   require Logger
   use Application
 
-  @impl true
-  def start(_type, _args) do
+  @upload_dir Application.app_dir(:app, ["priv", "static", "uploads"])
 
+  @saved_index if Application.compile_env(:app, :knnindex_indices_test, false),
+                 do: Path.join(@upload_dir, "indexes_test.bin"),
+                 else: Path.join(@upload_dir, "indexes.bin")
+
+  def check_models_on_startup do
     App.Models.verify_and_download_models()
+    |> case do
+      {:error, msg} ->
+        Logger.error("⚠️ #{msg}")
+        System.stop(0)
+
+      :ok ->
+        Logger.info("ℹ️ Models: ✅")
+        :ok
+    end
+  end
+
+  @impl true
+  def start(_type, _args) do
+    :ok = check_models_on_startup()
 
     children = [
       # Start the Telemetry supervisor
@@ -18,17 +36,16 @@ defmodule App.Application do
       # Start the PubSub system
       {Phoenix.PubSub, name: App.PubSub},
       # Nx serving for the embedding
-      # App.TextEmbedding,
-
+      {Nx.Serving, serving: App.Models.embedding(), name: Embedding, batch_size: 1},
       # Nx serving for Speech-to-Text
       {Nx.Serving,
-      serving:
-        if Application.get_env(:app, :use_test_models) == true do
-          App.Models.audio_serving_test()
-        else
-          App.Models.audio_serving()
-        end,
-      name: Whisper},
+       serving:
+         if Application.get_env(:app, :use_test_models) == true do
+           App.Models.audio_serving_test()
+         else
+           App.Models.audio_serving()
+         end,
+       name: Whisper},
       # Nx serving for image classifier
       {Nx.Serving,
        serving:
@@ -39,7 +56,7 @@ defmodule App.Application do
          end,
        name: ImageClassifier},
       {GenMagic.Server, name: :gen_magic},
-      
+
       # Adding a supervisor
       {Task.Supervisor, name: App.TaskSupervisor},
       # Start the Endpoint (http/https)
@@ -48,6 +65,17 @@ defmodule App.Application do
       # {App.Worker, arg}
     ]
 
+    # We are starting the HNSWLib Index GenServer only during testing.
+    # Because this GenServer needs the database to be seeded first,
+    # we only add it when we're not testing.
+    # When testing, you need to spawn this process manually (it is done in the test_helper.exs file).
+    children =
+      if Application.get_env(:app, :start_genserver, true) == true do
+        Enum.concat(children, [{App.KnnIndex, [space: :cosine, index: @saved_index]}])
+      else
+        children
+      end
+
     # See https://hexdocs.pm/elixir/Supervisor.html
     # for other strategies and supported options
     opts = [strategy: :one_for_one, name: App.Supervisor]