chore: Move vector precomputation into separate module

lib/medical_transcription/coding/vector_precomputation.ex

@@ -0,0 +1,80 @@
+defmodule MedicalTranscription.Coding.VectorPrecomputation do
+  @moduledoc "Populate database with vector embeddings from downloaded ICD-9 code list"
+
+  alias MedicalTranscription.CodeVector
+
+  @doc "Downloads the ICD-9 codelist, calculates vector embeddings for each, and adds them to the database"
+  def run() do
+    if File.exists?(csv_file()) do
+      IO.puts("CSV file found. Precomputing vectors...")
+    else
+      IO.puts("CSV file not found. Downloading and preparing...")
+      AudioTagger.SampleData.get_icd9_code_list_csv()
+
+      IO.puts("Precomputing vectors...")
+    end
+
+    precompute_vectors()
+
+    :ok
+  end
+
+  defp precompute_vectors() do
+    time_start = System.monotonic_time()
+    df = load_dataframe_from_csv()
+    model_tuple = AudioTagger.Classifier.SemanticSearch.prepare_model()
+
+    num_rows = Explorer.DataFrame.n_rows(df)
+    ProgressBar.render(0, num_rows, suffix: :count)
+
+    df
+    |> Explorer.DataFrame.to_rows_stream()
+    |> Stream.filter(fn %{"code" => code} -> String.length(code) > 0 end)
+    |> Stream.with_index()
+    |> Enum.each(fn {%{"code" => code, "long_description" => description}, index} ->
+      if !CodeVector.exists_for_code?(code) do
+        compute_vector_for_code(model_tuple, code, description)
+      end
+
+      ProgressBar.render(index + 1, num_rows, suffix: :count)
+    end)
+
+    ProgressBar.render(num_rows, num_rows, suffix: :count)
+    time_end = System.monotonic_time()
+
+    IO.puts(
+      "Finished in #{System.convert_time_unit(time_end - time_start, :native, :millisecond)}ms"
+    )
+  end
+
+  defp load_dataframe_from_csv() do
+    {:ok, df} =
+      Explorer.DataFrame.from_csv(
+        csv_file(),
+        dtypes: [
+          {"code", :string},
+          {"long_description", :string}
+        ]
+      )
+
+    df
+  end
+
+  defp compute_vector_for_code({model_info, tokenizer}, code, description) do
+    vector =
+      AudioTagger.Vectors.embed_with_model(model_info, tokenizer, [description])
+
+    vector_for_db = Nx.to_flat_list(vector.pooled_state)
+
+    CodeVector.insert_vector(%{
+      code: code,
+      description: description,
+      description_vector: vector_for_db
+    })
+  end
+
+  defp csv_file() do
+    AudioTagger.SampleData.cache_dir()
+    |> Path.join("icd9_codelist.csv")
+  end
+end

lib/mix/build_code_vectors.ex

@@ -2,83 +2,12 @@ defmodule Mix.Tasks.BuildCodeVectors do
   @moduledoc "Populate database with vector embeddings from downloaded ICD-9 code list"
 
   use Mix.Task
-  alias MedicalTranscription.CodeVector
 
   @shortdoc "Downloads the ICD-9 codelist, calculates vector embeddings for each, and adds them to the database"
   def run(_args) do
     Mix.Task.run("app.start")
     Logger.configure(level: :info)
 
-    if File.exists?(csv_file()) do
-      IO.puts("CSV file found. Precomputing vectors...")
-    else
-      IO.puts("CSV file not found. Downloading and preparing...")
-      AudioTagger.SampleData.get_icd9_code_list_csv()
-
-      IO.puts("Precomputing vectors...")
-    end
-
-    precompute_vectors()
-
-    :ok
-  end
-
-  defp precompute_vectors() do
-    time_start = System.monotonic_time()
-    df = load_dataframe_from_csv()
-    model_tuple = AudioTagger.Classifier.SemanticSearch.prepare_model()
-
-    num_rows = Explorer.DataFrame.n_rows(df)
-    ProgressBar.render(0, num_rows, suffix: :count)
-
-    df
-    |> Explorer.DataFrame.to_rows_stream()
-    |> Stream.filter(fn %{"code" => code} -> String.length(code) > 0 end)
-    |> Stream.with_index()
-    |> Enum.each(fn {%{"code" => code, "long_description" => description}, index} ->
-      if !CodeVector.exists_for_code?(code) do
-        compute_vector_for_code(model_tuple, code, description)
-      end
-
-      ProgressBar.render(index + 1, num_rows, suffix: :count)
-    end)
-
-    ProgressBar.render(num_rows, num_rows, suffix: :count)
-    time_end = System.monotonic_time()
-
-    IO.puts(
-      "Finished in #{System.convert_time_unit(time_end - time_start, :native, :millisecond)}ms"
-    )
-  end
-
-  defp load_dataframe_from_csv() do
-    {:ok, df} =
-      Explorer.DataFrame.from_csv(
-        csv_file(),
-        dtypes: [
-          {"code", :string},
-          {"long_description", :string}
-        ]
-      )
-
-    df
-  end
-
-  defp compute_vector_for_code({model_info, tokenizer}, code, description) do
-    vector =
-      AudioTagger.Vectors.embed_with_model(model_info, tokenizer, [description])
-
-    vector_for_db = Nx.to_flat_list(vector.pooled_state)
-
-    CodeVector.insert_vector(%{
-      code: code,
-      description: description,
-      description_vector: vector_for_db
-    })
-  end
-
-  defp csv_file() do
-    AudioTagger.SampleData.cache_dir()
-    |> Path.join("icd9_codelist.csv")
+    MedicalTranscription.Coding.VectorPrecomputation.run()
   end
 end