feat: Gather keywords from transcribed text

lib/medical_transcription/application.ex

@@ -14,6 +14,9 @@ defmodule MedicalTranscription.Application do
       |> AudioTagger.Transcriber.prepare_featurizer()
       |> AudioTagger.Transcriber.prepare_serving(@model_name)
 
+    token_classification_serving =
+      AudioTagger.KeywordFinder.prepare_token_classification_serving()
+
     children = [
       MedicalTranscriptionWeb.Telemetry,
       MedicalTranscription.Repo,
@@ -29,6 +32,13 @@ defmodule MedicalTranscription.Application do
         batch_size: 4,
         batch_timeout: 100
       },
+      {
+        Nx.Serving,
+        serving: token_classification_serving,
+        name: TokenClassificationServing,
+        batch_size: 1,
+        batch_timeout: 100
+      },
       # Start a worker by calling: MedicalTranscription.Worker.start_link(arg)
       # {MedicalTranscription.Worker, arg},
       # Start to serve requests, typically the last entry

lib/medical_transcription_web/components/components.ex

@@ -131,6 +131,7 @@ defmodule MedicalTranscriptionWeb.Components do
   end
 
   defp feedback_button(assigns) do
+    # TODO: Update UI after submitting feedback. E.g. change background color, etc. for confirmation.
     ~H"""
     <button
       phx-click="add_feedback"

lib/medical_transcription_web/components/layouts/app.html.heex

@@ -11,6 +11,7 @@
 
     <div class="px-6">
       <p class="text-xs leading-normal tracking-[0.2em] font-semibold uppercase">Today</p>
+      <!-- TODO: Show a history of uploaded files here. -->
     </div>
   </header>
 

lib/medical_transcription_web/live/home_live/index.ex

@@ -11,6 +11,7 @@ defmodule MedicalTranscriptionWeb.HomeLive.Index do
       |> assign(:uploaded_file_name, nil)
       |> stream(:transcription_rows, [])
       |> assign(:status, :pending)
+      |> assign(:keywords_for_rows, [])
       # To test the success UI, replace the two lines above with these:
       # |> stream(:transcription_rows, SampleResults.get_sample_results())
       # |> assign(:status, :success)
@@ -25,6 +26,7 @@ defmodule MedicalTranscriptionWeb.HomeLive.Index do
     # - Fine-tune results based on feedback:
     #   - Allow users to accept/decline suggested codes.
     #   - Train model based on user feedback for suggested codes.
+    # - After uploading an audio file, save the results to the DB (current date & time, transcription, assigned codes, etc.)
     # - Allow editing the transcription inline to correct mistakes. Then, retag based on the updated transcription.
     # - Stream audio recording instead of uploaded audio.
 
@@ -32,10 +34,27 @@ defmodule MedicalTranscriptionWeb.HomeLive.Index do
     <div class="flex-1 flex flex-col space-y-6">
       <%= if @status == :loading || @status == :success do %>
         <.result_heading status={@status} filename={@uploaded_file_name} />
+        <!-- TODO: Include this data inline with the transcription text. -->
+        <div>
+          <p>Summary Keywords:</p>
+          <ul>
+            <%= for keywords <- @keywords_for_rows do %>
+              <li class="list-disc">
+                <%= keywords
+                |> Enum.filter(fn keyword -> keyword.score > 0.25 end)
+                |> Enum.take(3)
+                |> Enum.map(fn keyword ->
+                  "#{keyword.label} (#{keyword.score})"
+                end)
+                |> Enum.join(", ") %>
+              </li>
+            <% end %>
+          </ul>
+        </div>
       <% end %>
 
       <%= if @status != :pending do %>
-        <.result_list rows={@streams.transcription_rows} />
+        <.result_list rows={@streams.transcription_rows} keywords_for_rows={@keywords_for_rows} />
       <% end %>
 
       <%= if @status == :pending do %>
@@ -77,6 +96,7 @@ defmodule MedicalTranscriptionWeb.HomeLive.Index do
       socket
       |> assign(:status, :loading)
       |> assign(:transcription_rows, [])
+      |> assign(:keywords_for_rows, [])
       |> assign(:uploaded_file_name, filename)
 
     {:noreply, socket}
@@ -109,14 +129,26 @@ defmodule MedicalTranscriptionWeb.HomeLive.Index do
   def handle_info({:transcription_row, chunk_result}, socket) do
     # The processing sends a message as each chunk of text is coded. See here for some background and potential
     # inspiration for this: https://elixirforum.com/t/liveview-asynchronous-task-patterns/44695
+
+    live_view_pid = self()
+    Task.async(fn -> find_keywords(live_view_pid, chunk_result) end)
+
     {:noreply, stream_insert(socket, :transcription_rows, chunk_result)}
   end
 
+  @impl true
+  def handle_info({:keywords_for_row, output}, socket) do
+    keywords = socket.assigns.keywords_for_rows
+
+    {:noreply, assign(socket, :keywords_for_rows, keywords ++ [output])}
+  end
+
   @impl true
   def handle_info({ref, _result}, socket) do
     # See this Fly article for the usage of Task.async to start `transcribe_and_tag_audio/2` and handle the end of the
     # task here: https://fly.io/phoenix-files/liveview-async-task/
     Process.demonitor(ref, [:flush])
+
     {:noreply, assign(socket, :status, :success)}
   end
 
@@ -127,6 +159,16 @@ defmodule MedicalTranscriptionWeb.HomeLive.Index do
     )
   end
 
+  defp find_keywords(live_view_pid, chunk) do
+    output = Nx.Serving.batched_run(TokenClassificationServing, chunk.text)
+    phrases = AudioTagger.KeywordFinder.cleanup_phrases(output.entities)
+
+    serving = AudioTagger.KeywordFinder.prepare_zero_shot_classification_serving(phrases)
+    output = Nx.Serving.run(serving, chunk.text)
+
+    send(live_view_pid, {:keywords_for_row, output.predictions})
+  end
+
   def error_to_string(:too_large), do: "Too large"
   def error_to_string(:not_accepted), do: "You have selected an unacceptable file type"
 end

livebooks/summarization.livemd

@@ -0,0 +1,157 @@
+# Text summarization
+
+```elixir
+Mix.install(
+  [
+    {:kino_bumblebee, "~> 0.4.0"},
+    {:exla, ">= 0.0.0"},
+    {:explorer, "~> 0.7.0"},
+    {:kino_explorer, "~> 0.1.11"}
+  ],
+  config: [nx: [default_backend: EXLA.Backend]]
+)
+```
+
+## Summarize an audio transcription
+
+```elixir
+text = """
+This 55-year-old man with known coronary artery disease comes for a follow-up visit today.
+Last month he was admitted to our hospital with unstable angina.
+He underwent heart catheterization on November 15th, 2007.
+At that time he was found to have a tight 99% proxmost enosis, total occlusion and collateralization
+of the mid-circumflex, right coronary artery was normal.
+Ventricularography was normal and his ejection fraction was 65%.
+He underwants an uncomplicated placement of a cipher drug-eleuting
+stent to his proximal lesion.
+The attempted coronary intervention of the circumflex was unsuccessful,
+as his lesion cannot be crossed. His post procedure was uncomplicated, and he was discharged on the day following his intervention.
+He comes today indicating that he is feeling great.
+His current medications include aspirin,
+325 milligrams daily, lipatore, 40 milligrams daily,
+and platvic, 75 milligrams daily.
+"""
+```
+
+## Section
+
+```elixir
+{:ok, model_info} =
+  Bumblebee.load_model({:hf, "vblagoje/bert-english-uncased-finetuned-pos"})
+
+{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "bert-base-uncased"})
+
+serving =
+  Bumblebee.Text.token_classification(model_info, tokenizer,
+    aggregation: :same,
+    compile: [batch_size: 1, sequence_length: 100],
+    defn_options: [compiler: EXLA]
+  )
+```
+
+```elixir
+text = "This 55-year-old man with known coronary artery disease comes for a follow-up visit today.
+Last month he was admitted to our hospital with unstable angina."
+
+text = "Last month he was admitted to our hospital with unstable angina."
+
+ignored = ["DET", "PUNCT", "ADP", "NUM", "AUX", "PRON"]
+# ignored = []
+
+output = Nx.Serving.run(serving, text)
+
+output.entities
+|> Enum.reduce([], fn entity, acc ->
+  if Enum.member?(ignored, entity.label) do
+    acc
+  else
+    # "VERB", "NOUN", and "ADJ"
+    next_phrase =
+      if entity.label == "ADJ" do
+        "#{entity.phrase} [CONTINUATION]"
+      else
+        entity.phrase
+      end
+
+    if Enum.count(acc) > 0 do
+      previous = Enum.at(acc, -1)
+
+      # First, check if the previous phrase ends with a continuation token.  
+      if String.ends_with?(previous, "[CONTINUATION]") do
+        acc_without_last = Enum.take(acc, Enum.count(acc) - 1)
+        acc_without_last ++ [String.replace(previous, "[CONTINUATION]", next_phrase)]
+      else
+        acc ++ [next_phrase]
+      end
+    else
+      acc ++ [next_phrase]
+    end
+  end
+end)
+```
+
+<!-- livebook:{"attrs":"eyJjb21waWxlciI6ImV4bGEiLCJsYWJlbHMiOiJsYXN0IG1vbnRoLCBhZG1pdHRlZCwgaG9zcGl0YWwsIHVuc3RhYmxlIGFuZ2luYSIsInNlcXVlbmNlX2xlbmd0aCI6MTAwLCJ0YXNrX2lkIjoiemVyb19zaG90X3RleHRfY2xhc3NpZmljYXRpb24iLCJ0b3BfayI6bnVsbCwidmFyaWFudF9pZCI6ImJhcnRfbGFyZ2VfbW5saSJ9","chunks":[[0,396],[398,509]],"kind":"Elixir.KinoBumblebee.TaskCell","livebook_object":"smart_cell"} -->
+
+```elixir
+{:ok, model_info} = Bumblebee.load_model({:hf, "facebook/bart-large-mnli"})
+{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "facebook/bart-large-mnli"})
+labels = ["last month", "admitted", "hospital", "unstable angina"]
+
+serving =
+  Bumblebee.Text.zero_shot_classification(model_info, tokenizer, labels,
+    compile: [batch_size: 1, sequence_length: 100],
+    defn_options: [compiler: EXLA]
+  )
+
+text_input = Kino.Input.textarea("Text", default: "One day I will see the world.")
+form = Kino.Control.form([text: text_input], submit: "Run")
+frame = Kino.Frame.new()
+
+Kino.listen(form, fn %{data: %{text: text}} ->
+  Kino.Frame.render(frame, Kino.Text.new("Running..."))
+  output = Nx.Serving.run(serving, text)
+
+  output.predictions
+  |> Enum.map(&{&1.label, &1.score})
+  |> Kino.Bumblebee.ScoredList.new()
+  |> then(&Kino.Frame.render(frame, &1))
+end)
+
+Kino.Layout.grid([form, frame], boxed: true, gap: 16)
+```
+
+## Question answering
+
+<!-- livebook:{"attrs":"eyJjb21waWxlciI6ImV4bGEiLCJzZXF1ZW5jZV9sZW5ndGgiOjUwMCwidGFza19pZCI6InF1ZXN0aW9uX2Fuc3dlcmluZyIsInZhcmlhbnRfaWQiOiJkaXN0aWxiZXJ0X2Jhc2VfY2FzZWQifQ","chunks":[[0,344],[346,595]],"kind":"Elixir.KinoBumblebee.TaskCell","livebook_object":"smart_cell"} -->
+
+```elixir
+{:ok, model_info} = Bumblebee.load_model({:hf, "distilbert-base-cased-distilled-squad"})
+
+{:ok, tokenizer} =
+  Bumblebee.load_tokenizer({:hf, "distilbert-base-cased-distilled-squad"})
+
+serving =
+  Bumblebee.Text.question_answering(model_info, tokenizer,
+    compile: [batch_size: 1, sequence_length: 500],
+    defn_options: [compiler: EXLA]
+  )
+
+inputs = [
+  question: Kino.Input.text("Question", default: "Where do I live?"),
+  context: Kino.Input.textarea("Context", default: "My name is Sarah and I live in London.")
+]
+
+form = Kino.Control.form(inputs, submit: "Run")
+frame = Kino.Frame.new()
+
+Kino.listen(form, fn %{data: %{question: question, context: context}} ->
+  output = Nx.Serving.run(serving, %{question: question, context: context})
+
+  output.results
+  |> Enum.map(&{&1.text, &1.score})
+  |> Kino.Bumblebee.ScoredList.new()
+  |> then(&Kino.Frame.render(frame, &1))
+end)
+
+Kino.Layout.grid([form, frame], boxed: true, gap: 16)
+```