feat: Add :progress_bar_step config option (#441)

Author: georgeguimaraes

README.md

@@ -129,6 +129,18 @@ First, if the repository is clearly a fine-tuned version of another model, you c
 
 Otherwise, the Transformers library includes conversion rules to load a "slow tokenizer" and convert it to a corresponding "fast tokenizer", which is possible in most cases. You can generate the `tokenizer.json` file using [this tool](https://jonatanklosko-bumblebee-tools.hf.space/apps/tokenizer-generator). Once successful, you can follow the steps to submit a PR adding `tokenizer.json` to the model repository. Note that you do not have to wait for the PR to be merged, instead you can copy commit SHA from the PR and load the tokenizer with `Bumblebee.load_tokenizer({:hf, "model-repo", revision: "..."})`.
 
+## Global configuration
+
+You can configure the progress bar displayed during model downloads:
+
+```elixir
+# Update every 10% instead of every 1%
+config :bumblebee, :progress_bar_step, 10
+
+# Disable progress bar entirely
+config :bumblebee, :progress_bar_enabled, false
+```
+
 <!-- Docs -->
 
 ## Contributing

lib/bumblebee/utils.ex

@@ -8,4 +8,15 @@ defmodule Bumblebee.Utils do
   def progress_bar_enabled?() do
     Application.get_env(:bumblebee, :progress_bar_enabled, true)
   end
+
+  @doc """
+  Returns the progress bar update step in percent.
+
+  Progress updates only when crossing step boundaries (e.g., every 1%).
+  Defaults to `1`. Set to `nil` for updates on every chunk.
+  """
+  @spec progress_bar_step :: non_neg_integer() | nil
+  def progress_bar_step() do
+    Application.get_env(:bumblebee, :progress_bar_step, 1)
+  end
 end

lib/bumblebee/utils/http.ex

@@ -49,7 +49,13 @@ defmodule Bumblebee.Utils.HTTP do
           opts = [stream: :self, sync: false, receiver: receiver]
 
           {:ok, request_id} = :httpc.request(:get, request, http_opts, opts, :bumblebee)
-          download_loop(%{request_id: request_id, file: file, total_size: nil, size: nil})
+
+          download_loop(%{
+            request_id: request_id,
+            file: file,
+            total_size: nil,
+            size: nil
+          })
         after
           File.close(file)
         end
@@ -88,12 +94,7 @@ defmodule Bumblebee.Utils.HTTP do
       :ok ->
         part_size = byte_size(body_part)
         state = update_in(state.size, &(&1 + part_size))
-
-        if Bumblebee.Utils.progress_bar_enabled?() &&
-             state.total_size && part_size != state.total_size do
-          ProgressBar.render(state.size, state.total_size, suffix: :bytes)
-        end
-
+        state = maybe_render_progress(state, part_size)
         download_loop(state)
 
       {:error, error} ->
@@ -106,6 +107,27 @@ defmodule Bumblebee.Utils.HTTP do
     :ok
   end
 
+  defp maybe_render_progress(state, part_size)
+       when is_nil(state.total_size) or part_size == state.total_size do
+    state
+  end
+
+  defp maybe_render_progress(state, part_size) do
+    if Bumblebee.Utils.progress_bar_enabled?() do
+      step = Bumblebee.Utils.progress_bar_step()
+      percent = trunc(state.size / state.total_size * 100)
+      last_percent = trunc((state.size - part_size) / state.total_size * 100)
+      step_bucket = if step, do: div(percent, step), else: percent
+      last_step_bucket = if step, do: div(last_percent, step), else: last_percent
+
+      if step_bucket > last_step_bucket or percent == 100 do
+        ProgressBar.render(state.size, state.total_size, suffix: :bytes)
+      end
+    end
+
+    state
+  end
+
   defp total_size(headers) do
     case List.keyfind(headers, ~c"content-length", 0) do
       {_, content_length} ->