feat: Identify and store chunk keywords

lib/medical_transcription/application.ex

@@ -23,9 +23,11 @@ defmodule MedicalTranscription.Application do
       text_embedding_spec(),
       {
         MedicalTranscription.TranscriptionSupervisor,
-        # name: MedicalTranscription.TranscriptionSupervisor,
-        strategy: :one_for_one,
-        max_restarts: 1
+        strategy: :one_for_one, max_restarts: 1
+      },
+      {
+        MedicalTranscription.ClassificationSupervisor,
+        strategy: :one_for_one, max_restarts: 1
       },
       # Start a worker by calling: MedicalTranscription.Worker.start_link(arg)
       # {MedicalTranscription.Worker, arg},

lib/medical_transcription/classification_server.ex

@@ -0,0 +1,108 @@
+defmodule MedicalTranscription.ClassificationServer do
+  @moduledoc """
+  GenServer responsible for classifying transcription text
+  """
+  use GenServer
+
+  alias AudioTagger.KeywordFinder
+
+  alias MedicalTranscription.Transcriptions
+  alias MedicalTranscription.Transcriptions.TranscriptionChunk
+
+  def start_link(args) do
+    GenServer.start_link(__MODULE__, args, [])
+  end
+
+  @impl GenServer
+  def init(init_arg) do
+    {:ok, init_arg, {:continue, :start}}
+  end
+
+  @impl GenServer
+  def handle_continue(:start, {:chunk, chunk} = state) do
+    find_keywords(chunk)
+
+    Phoenix.PubSub.broadcast(
+      :medicode_pubsub,
+      "transcription_chunks:#{chunk.id}",
+      {:classification_started, chunk.id}
+    )
+
+    {:noreply, state}
+  end
+
+  @impl GenServer
+  def handle_info({:chunk_updated, result}, state) do
+    {:chunk, chunk} = state
+
+    %TranscriptionChunk{id: id} = chunk
+
+    Phoenix.PubSub.broadcast(
+      :medicode_pubsub,
+      "transcription_chunks:#{id}",
+      {:chunk_updated, chunk}
+    )
+
+    {:noreply, state}
+  end
+
+  def handle_info(:finished, state) do
+    {:stop, :normal, state}
+  end
+
+  @impl GenServer
+  def terminate(reason, state) do
+    {:chunk, chunk} = state
+
+    %TranscriptionChunk{id: id} = chunk
+
+    Phoenix.PubSub.broadcast(
+      :medicode_pubsub,
+      "transcription_chunks:#{id}",
+      {:classification_finished, reason}
+    )
+
+    reason
+  end
+
+  defp classify_chunk(chunk) do
+    find_keywords(chunk)
+  end
+
+  defp find_keywords(chunk) do
+    %{entities: entities} =
+      Nx.Serving.batched_run(MedicalTranscription.TokenClassificationServing, chunk.text)
+
+    phrases = KeywordFinder.cleanup_phrases(entities)
+
+    # Then, we use one of two processes to determine which to show as keywords
+    keywords =
+      chunk.text
+      |> determine_keywords(phrases)
+      |> Enum.map(fn %{label: label, score: score} ->
+        # TODO: Replace loop with an insert_all call and check for conflicts
+        # so that duplicate keywords are ignored.
+        Transcriptions.create_keyword_for_chunk(%{
+          transcription_chunk_id: chunk.id,
+          keyword: label,
+          score: score
+        })
+      end)
+
+    send(self(), {:chunk_updated, keywords})
+
+    send(self(), :finished)
+  end
+
+  # This clause handles cases where there is transcribed text, but no phrases were found.
+  defp determine_keywords(_text, []), do: []
+
+  defp determine_keywords(text, phrases) do
+    # 1. A slower process that looks to classify the text by the extracted phrases.
+    # serving = KeywordFinder.prepare_zero_shot_classification_serving(phrases)
+    # %{predictions: predictions} = Nx.Serving.run(serving, text)
+
+    # 2. A fast process finding the phrase closest in vector space to the whole text.
+    KeywordFinder.find_most_similar_label(text, phrases, 2)
+  end
+end

lib/medical_transcription/classification_supervisor.ex

@@ -0,0 +1,25 @@
+defmodule MedicalTranscription.ClassificationSupervisor do
+  # Automatically defines child_spec/1
+  use DynamicSupervisor
+
+  def start_link(init_arg) do
+    DynamicSupervisor.start_link(__MODULE__, init_arg, name: __MODULE__)
+  end
+
+  @impl true
+  def init(_init_arg) do
+    DynamicSupervisor.init(strategy: :one_for_one)
+  end
+
+  def start_classification(transcription_chunk) do
+    spec = %{
+      id: MedicalTranscription.ClassificationServer,
+      start:
+        {MedicalTranscription.ClassificationServer, :start_link, [chunk: transcription_chunk]},
+      restart: :transient,
+      type: :worker
+    }
+
+    DynamicSupervisor.start_child(__MODULE__, spec)
+  end
+end

lib/medical_transcription/transcriptions.ex

@@ -57,7 +57,7 @@ defmodule MedicalTranscription.Transcriptions do
     query =
       if preload_transcription_chunks do
         Transcription
-        |> preload(:chunks)
+        |> preload(chunks: :keywords)
       else
         Transcription
       end
@@ -83,7 +83,7 @@ defmodule MedicalTranscription.Transcriptions do
     query =
       if preload_transcription_chunks do
         Transcription
-        |> preload(:chunks)
+        |> preload(chunks: :keywords)
       else
         Transcription
       end
@@ -127,6 +127,24 @@ defmodule MedicalTranscription.Transcriptions do
     |> Repo.insert()
   end
 
+  @doc """
+  Creates a transcription chunk keyword.
+
+  ## Examples
+
+      iex> create_keyword_for_chunk(%{transcription_chunk_id: 1, keyword: "healthy"})
+      {:ok, %TranscriptionChunkKeyword{}}
+
+      iex> create_keyword_for_chunk(%{field: bad_value})
+      {:error, %Ecto.Changeset{}}
+
+  """
+  def create_keyword_for_chunk(attrs \\ %{}) do
+    %TranscriptionChunkKeyword{}
+    |> TranscriptionChunkKeyword.changeset(attrs)
+    |> Repo.insert()
+  end
+
   @doc """
   Updates a transcription.
 

lib/medical_transcription/transcriptions/transcription_chunk.ex

@@ -11,6 +11,8 @@ defmodule MedicalTranscription.Transcriptions.TranscriptionChunk do
 
     belongs_to :transcription, MedicalTranscription.Transcriptions.Transcription
 
+    has_many :keywords, MedicalTranscription.Transcriptions.TranscriptionChunkKeyword
+
     timestamps(type: :utc_datetime)
   end
 

lib/medical_transcription/transcriptions/transcription_chunk_keyword.ex

@@ -0,0 +1,22 @@
+defmodule MedicalTranscription.Transcriptions.TranscriptionChunkKeyword do
+  use Ecto.Schema
+  import Ecto.Changeset
+
+  @primary_key {:id, :binary_id, autogenerate: true}
+  @foreign_key_type :binary_id
+  schema "transcription_chunk_keywords" do
+    field :keyword, :string
+    field :score, :float
+
+    belongs_to :transcription_chunk, MedicalTranscription.Transcriptions.TranscriptionChunk
+
+    timestamps(type: :utc_datetime)
+  end
+
+  @doc false
+  def changeset(transcription_chunk_keyword, attrs) do
+    transcription_chunk_keyword
+    |> cast(attrs, [:transcription_chunk_id, :keyword, :score])
+    |> validate_required([:transcription_chunk_id, :keyword, :score])
+  end
+end

priv/repo/migrations/20240209165222_create_transcription_chunk_keywords.exs

@@ -0,0 +1,17 @@
+defmodule MedicalTranscription.Repo.Migrations.CreateTranscriptionChunkKeywords do
+  use Ecto.Migration
+
+  def change do
+    create table(:transcription_chunk_keywords, primary_key: false) do
+      add :id, :binary_id, primary_key: true
+      add :keyword, :string
+      add :score, :float
+
+      add :transcription_chunk_id,
+          references(:transcription_chunks, type: :binary_id, on_delete: :delete_all),
+          null: false
+
+      timestamps(type: :utc_datetime)
+    end
+  end
+end

test/medical_transcription/classification_server_test.exs

@@ -0,0 +1,42 @@
+defmodule MedicalTranscription.ClassificationServerTest do
+  @moduledoc """
+  Tests for MedicalTranscription.ClassificationServer
+  """
+
+  use MedicalTranscription.DataCase
+
+  import MedicalTranscription.TranscriptionsFixtures
+  import MedicalTranscription.TranscriptionChunksFixtures
+
+  alias MedicalTranscription.Transcriptions
+  alias MedicalTranscription.ClassificationServer
+
+  setup do
+    transcription = transcription_fixture()
+
+    chunk =
+      transcription_chunk_fixture(%{
+        transcription_id: transcription.id,
+        text: "patient is suffering from lower back pain"
+      })
+
+    %{chunk: chunk, transcription: transcription}
+  end
+
+  test "extract keywords and assign medical codes", %{chunk: chunk, transcription: transcription} do
+    spec = {ClassificationServer, {:chunk, chunk}}
+
+    {:ok, pid} = start_supervised(spec, restart: :transient)
+
+    ref = Process.monitor(pid)
+    assert_receive({:DOWN, ^ref, :process, _object, _pid}, 5_000)
+
+    keywords =
+      transcription.id
+      |> Transcriptions.get_transcription!(true)
+      |> Map.fetch!(:chunks)
+      |> Enum.flat_map(& &1.keywords)
+
+    assert 2 == Enum.count(keywords)
+  end
+end

test/medical_transcription/classification_supervisor_test.exs

@@ -0,0 +1,28 @@
+defmodule MedicalTranscription.ClassificationSupervisorTest do
+  @moduledoc """
+  Tests for MedicalTranscription.ClassificationSupervisor
+  """
+
+  use MedicalTranscription.DataCase
+
+  import MedicalTranscription.TranscriptionChunksFixtures
+
+  alias MedicalTranscription.ClassificationSupervisor
+
+  setup do
+    chunk = transcription_chunk_fixture(%{text: "patient fell and is experiecing knee pain"})
+    %{chunk: chunk}
+  end
+
+  test "extract keywords and assign codes", %{chunk: chunk} do
+    assert {:ok, _pid} = ClassificationSupervisor.start_classification(chunk)
+
+    {_, server_pid, :worker, [MedicalTranscription.ClassificationServer]} =
+      ClassificationSupervisor
+      |> DynamicSupervisor.which_children()
+      |> Enum.at(0)
+
+    ref = Process.monitor(server_pid)
+    assert_receive({:DOWN, ^ref, :process, _object, _pid}, 5_000)
+  end
+end