whisper server elixir

whisper_server/lib/whisper_server/application.ex

@@ -0,0 +1,55 @@
+defmodule WhisperServer.Application do
+  use Application
+
+  def start(_type, _args) do
+    args = parse_args(System.argv())
+
+    Application.put_env(:whisper_server, :model_name, args[:model])
+    Application.put_env(:whisper_server, :client, String.to_atom(args[:client]))
+    Application.put_env(:whisper_server, :batch_size, args[:batch_size])
+    Application.put_env(:whisper_server, :batch_timeout, args[:batch_timeout])
+    Application.put_env(:whisper_server, :port, args[:port])
+
+    children = [
+      WhisperServer.WhisperInference,
+      {Plug.Cowboy, scheme: :http, plug: WhisperServer, options: [port: args[:port]]}
+    ]
+
+    opts = [strategy: :one_for_one, name: WhisperServer.Supervisor]
+    Supervisor.start_link(children, opts)
+  end
+
+  defp parse_args(argv) do
+    OptionParser.parse!(argv,
+      switches: [
+        batch_size: :integer,
+        batch_timeout: :integer,
+        client: :string,
+        model: :string,
+        port: :integer
+      ],
+      aliases: [
+        b: :batch_size,
+        t: :batch_timeout,
+        c: :client,
+        m: :model,
+        p: :port
+      ]
+    )
+    |> elem(0)
+    # |> Enum.into(%{
+    #   batch_size: System.get_env("BATCH_SIZE") || 3,                
+    #   batch_timeout: System.get_env("BATCH_TIMEOUT") || 3000,        
+    #   client: System.get_env("CLIENT") || "host",               
+    #   model: System.get_env("MODEL") || "openai/whisper-tiny", 
+    #   port: System.get_env("PORT") || 4000
+    # })
+    |> Enum.into(%{
+      batch_size: String.to_integer(System.get_env("BATCH_SIZE") || "3"),
+      batch_timeout: String.to_integer(System.get_env("BATCH_TIMEOUT") || "3000"),
+      client: System.get_env("CLIENT") || "host",
+      model: System.get_env("MODEL") || "openai/whisper-tiny",
+      port: String.to_integer(System.get_env("PORT") || "4000")
+    })
+  end
+end

whisper_server/lib/whisper_server/inference_runner.ex

@@ -0,0 +1,10 @@
+defmodule WhisperServer.InferenceRunner do
+  @moduledoc """
+  Runs inference on audio files using the initialized Whisper model.
+  """
+
+  def run_inference(audio_path) do
+    result = Nx.Serving.batched_run(WhisperServer.WhisperInference.Serving, {:file, audio_path})
+    result
+  end
+end

whisper_server/lib/whisper_server/whisper_inference.ex

@@ -0,0 +1,51 @@
+defmodule WhisperServer.WhisperInference do
+  use Supervisor
+
+  @moduledoc """
+  Initializes the Whisper model and sets up the serving process.
+  """
+
+  def start_link(opts) do
+    Supervisor.start_link(__MODULE__, opts, name: __MODULE__)
+  end
+
+  def init(_opts) do
+    model_name = Application.get_env(:whisper_server, :model_name, "openai/whisper-tiny")
+    raw_client = Application.get_env(:whisper_server, :client, :host)
+
+    client =
+      case raw_client do
+        :rocm ->
+          IO.warn("Client :rocm is not supported, falling back to :host")
+          :host
+        _ -> raw_client
+      end
+
+    batch_size = Application.get_env(:whisper_server, :batch_size, 3)
+    batch_timeout = Application.get_env(:whisper_server, :batch_timeout, 3000)
+
+    Nx.global_default_backend({EXLA.Backend, client: client})
+
+    {:ok, model} = Bumblebee.load_model({:hf, model_name})
+    {:ok, featurizer} = Bumblebee.load_featurizer({:hf, model_name})
+    {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, model_name})
+    {:ok, generation_config} = Bumblebee.load_generation_config({:hf, model_name})
+
+    serving = Bumblebee.Audio.speech_to_text_whisper(
+      model, featurizer, tokenizer, generation_config,
+      chunk_num_seconds: 30,
+      language: "es",
+      defn_options: [compiler: EXLA, client: :host]
+    )
+
+    children = [
+      {Nx.Serving,
+       serving: serving,
+       name: __MODULE__.Serving,
+       batch_size: batch_size,
+       batch_timeout: batch_timeout}
+    ]
+
+    Supervisor.init(children, strategy: :one_for_one)
+  end
+end

whisper_server/lib/whisper_server/whisper_server.ex

@@ -0,0 +1,127 @@
+defmodule WhisperServer do
+  use Plug.Router
+  require Logger
+
+  plug Plug.Parsers,
+    parsers: [:urlencoded, :multipart, :json],
+    pass: ["*/*"],
+    json_decoder: Jason
+
+  plug :match
+  plug :dispatch
+
+  post "/infer" do
+    handle_request(conn)
+  end
+
+  post "/v1/audio/transcriptions" do
+    model = conn.params["model"] || "whisper-1"
+    response_format = conn.params["response_format"] || "json"
+
+    if model != "whisper-1" do
+      send_resp(conn, 400, Jason.encode!(%{error: "Unsupported model"}))
+    else
+      upload = conn.params["file"]
+
+      case File.read(upload.path) do
+        {:ok, file_bytes} ->
+
+          filename = "uploaded_#{System.unique_integer([:positive])}_#{upload.filename}"
+          temp_path = Path.join("uploads", filename)
+
+          File.mkdir_p!("uploads")
+
+          case File.write(temp_path, file_bytes) do
+            :ok ->
+              try do
+                result = WhisperServer.InferenceRunner.run_inference(temp_path)
+                Logger.info("Inference result: #{inspect(result)}")
+                result_text = extract_text_from_infer_response(result)
+                Logger.info("Extracted text: #{result_text}")
+
+                case response_format do
+                  "text" ->
+                    conn
+                    |> put_resp_header("Content-Disposition", "attachment; filename=result.txt")
+                    |> send_resp(200, result_text)
+
+                  "json" ->
+                    conn
+                    |> put_resp_header("Content-Disposition", "attachment; filename=result.json")
+                    |> send_resp(200, Jason.encode!(%{text: result_text}))
+
+                  _ ->
+                    send_resp(conn, 200, Jason.encode!(result))
+                end
+              after
+                File.rm(temp_path)
+              end
+
+            {:error, reason} ->
+              send_resp(conn, 500, Jason.encode!(%{error: "Failed to save file: #{reason}"}))
+          end
+
+        {:error, reason} ->
+          send_resp(conn, 500, Jason.encode!(%{error: "Failed to read file: #{reason}"}))
+      end
+    end
+  end
+
+  post "/v1/audio/translations" do
+    send_resp(conn, 200, Jason.encode!(%{}))
+  end
+
+  get "/health" do
+    send_resp(conn, 200, Jason.encode!(%{status: "ok"}))
+  end
+
+  get "/v1/models" do
+    send_resp(conn, 200, Jason.encode!(["whisper-1"]))
+  end
+
+  get "/v1/models/:model" do
+    model = conn.params["model"]
+    if model == "whisper-1" do
+      send_resp(conn, 200, Jason.encode!(%{name: "whisper-1"}))
+    else
+      send_resp(conn, 404, Jason.encode!(%{error: "Model not found"}))
+    end
+  end
+
+  match _ do
+    send_resp(conn, 404, "Not Found")
+  end
+
+  defp extract_text_from_infer_response(response) do
+    response
+    |> Map.get(:chunks, [])
+    |> Enum.map(& &1[:text])
+    |> Enum.join(" ")
+    |> String.trim()
+  end
+
+  defp handle_request(conn) do
+    upload = conn.params["file"]
+
+    temp_path = decode_audio_from_body(upload)
+
+    try do
+      result = WhisperServer.InferenceRunner.run_inference(temp_path)
+
+      send_resp(conn, 200, Jason.encode!(result))
+    after
+      File.rm(temp_path)
+    end
+  end
+
+  defp decode_audio_from_body(%Plug.Upload{path: uploaded_file_path, filename: filename}) do
+    unique_name = "uploaded_#{System.unique_integer([:positive])}_#{filename}"
+    temp_path = Path.join("uploads", unique_name)
+
+    File.mkdir_p!("uploads")
+
+    File.cp!(uploaded_file_path, temp_path)
+
+    temp_path
+  end
+end