elixir and bumblebee with turbo v3 large

turbo.elixir

Mix.install([
  {:axon, "~> 0.7"},
  {:bumblebee, "~> 0.6"},
  {:exla, "~> 0.9"},
  {:nx, "~> 0.9"},
  {:mp3_duration, "~> 0.1.0"},
  {:plug_cowboy, "~> 2.6"},
  {:jason, "~> 1.4"},
  {:phoenix, "~> 1.7.14"},
  {:phoenix_html, "~> 4.1"},
  {:phoenix_live_view, "~> 1.0.0-rc.1", override: true}
])

Application.put_env(:phoenix, :json_library, Jason)

Application.put_env(:phoenix_demo, PhoenixDemo.Endpoint,
  http: [ip: {0, 0, 0, 0}, port: 4200],
  check_origin: false,
  server: true,
  live_view: [signing_salt: :crypto.strong_rand_bytes(8) |> Base.encode16()],
  secret_key_base: :crypto.strong_rand_bytes(32) |> Base.encode16(),
  pubsub_server: PhoenixDemo.PubSub
)

defmodule PhoenixDemo.Layouts do
  use Phoenix.Component

  def render("live.html", assigns) do
    ~H"""
    <script src="//cdn.jsdelivr.net/npm/[email protected]/priv/static/phoenix.min.js"></script>
    <script src="//cdn.jsdelivr.net/npm/[email protected]/priv/static/phoenix_live_view.min.js"></script>
    <script>
      const liveSocket = new LiveView.LiveSocket("/live", Phoenix.Socket, {hooks: {}})
      liveSocket.connect()
    </script>
    <script src="https://cdn.tailwindcss.com"></script>
    <%= @inner_content %>
    """
  end
end

defmodule PhoenixDemo.ErrorView do
  def render(_, _), do: "error"
end

defmodule PhoenixDemo.SampleLive do
  use Phoenix.LiveView, layout: {PhoenixDemo.Layouts, :live}

  @impl true
  def mount(params, _, socket) do
    socket =
      socket
      |> assign(task: nil, showdoc: true)
      |> allow_upload(:document, accept: ~w(.mp3), progress: &handle_progress/3, auto_upload: true, max_entries: 1, max_file_size: 100_000_000)
      |> stream_configure(:segments, dom_id: &"texts-#{&1.id}")
      |> stream(:segments, [])

    {:ok, socket}
  end

  @impl true
  def handle_event("noop", %{}, socket) do
    # We need phx-change and phx-submit on the form for live uploads
    {:noreply, socket}
  end

  @impl true
  def handle_info({ref, results}, socket) when socket.assigns.task.ref == ref do
    socket =
      results
      |> Enum.with_index()
      |> Enum.reduce(socket, fn {{_duration, ss, text}, index}, socket ->
        socket |> stream_insert(:segments, %{id: ss, text: text}, at: index)
      end)

    socket = socket |> assign(task: nil)

    {:noreply, socket}
  end

  @impl true
  def handle_info(_, socket) do
    {:noreply, socket}
  end

  def handle_progress(:document, %{client_name: filename} = entry, socket) when entry.done? do
    filepath =
      consume_uploaded_entry(socket, entry, fn %{path: path} ->
        :ok = File.cp!(path, "example.mp3")
        {:ok, "example.mp3"}
      end)

    {:ok, %{duration: duration}} = Mp3Duration.parse(filepath)

    task =
      speech_to_text(duration, filepath, 20, fn ss, text ->
        {duration, ss, text}
      end)

    {:noreply, assign(socket, task: task, showdoc: false)}
  end

  def handle_progress(_name, _entry, socket), do: {:noreply, socket}

  def speech_to_text(duration, path, chunk_time, func) do
    Task.async(fn ->
      format = get_format()

      0..duration//chunk_time
      |> Task.async_stream(
        fn ss ->
          args = ~w(-ac 1 -ar 16k -f #{format} -ss #{ss} -t #{chunk_time} -v quiet -)
          {data, 0} = System.cmd("ffmpeg", ["-i", path] ++ args)

          if is_nil(data) || data == "" do
            {ss, %{results: []}}
          else
            {ss, Nx.Serving.batched_run(WhisperServing, Nx.from_binary(data, :f32))}
          end
        end,
        max_concurrency: 4,
        timeout: :infinity
      )
      |> Enum.reject(fn {:ok, {_ss, %{chunks: results}}} -> Enum.empty?(results) end)
      |> Enum.map(fn {:ok, {ss, %{chunks: [%{text: text}]}}} ->
        func.(ss, text)
      end)
    end)
  end

  def get_format() do
    case System.endianness() do
      :little -> "f32le"
      :big -> "f32be"
    end
  end

  @impl true
  def render(assigns) do
    ~H"""
    <div class="h-screen px-8 py-4">
      <div id="segments" phx-update="stream" class="pt-4">
        <div
          id={id}
          :for={{id, segment} <- @streams.segments}
          class="flex w-full justify-center items-center text-blue-400 font-bold"
        >
          <%= segment.text %>
        </div>
      </div>
      <div :if={@showdoc} class="flex">
        <form phx-change="noop" phx-submit="noop" phx-drop-target={@uploads.document.ref}>
          <.live_file_input class="demospot" upload={@uploads.document} />
        </form>
      </div>
    </div>
    """
  end


end

defmodule PhoenixDemo.Router do
  use Phoenix.Router
  import Phoenix.LiveView.Router

  pipeline :browser do
    plug(:accepts, ["html"])
  end

  scope "/", PhoenixDemo do
    pipe_through(:browser)

    live("/", SampleLive, :index)
  end
end

defmodule PhoenixDemo.Endpoint do
  use Phoenix.Endpoint, otp_app: :phoenix_demo

  socket("/live", Phoenix.LiveView.Socket)
  plug(PhoenixDemo.Router)
end

# Application startup
Nx.default_backend(EXLA.Backend)

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

repository_id = {:hf, "openai/whisper-large-v3-turbo"}

{:ok, model_info} = Bumblebee.load_model(repository_id, type: :f16)
{:ok, featurizer} = Bumblebee.load_featurizer(repository_id)
{:ok, tokenizer} = Bumblebee.load_tokenizer(repository_id)

{:ok, generation_config} =
  Bumblebee.load_generation_config(repository_id)

serving =
  Bumblebee.Audio.speech_to_text_whisper(
    model_info,
    featurizer,
    tokenizer,
    generation_config,
    compile: [batch_size: 1],
    defn_options: [compiler: EXLA]
  )

{:ok, _} =
  Supervisor.start_link(
    [
      {Phoenix.PubSub, name: PhoenixDemo.PubSub},
      {Nx.Serving, serving: serving, name: WhisperServing},
      PhoenixDemo.Endpoint
    ],
    strategy: :one_for_one
  )

Process.sleep(:infinity)

if anyone wants to deploy this with fly here is a Dockerfile that works with Nx 0.9, bumblebee 0.6, cudnn 9 and cuda 12

https://github.com/toranb/mistral-chat-f16/blob/main/Dockerfile