From 922e1508c5de4412df02eeaaa3b6fdae32f9b6d3 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Thu, 1 Oct 2026 20:28:14 +0200 Subject: [PATCH 1/3] Support recursive repository file discovery and options for local repos --- lib/bumblebee.ex | 87 +++++++++++++++++++++---- lib/bumblebee/huggingface/hub.ex | 12 +++- test/bumblebee/huggingface/hub_test.exs | 18 +++++ test/bumblebee_test.exs | 19 ++++++ 4 files changed, 120 insertions(+), 16 deletions(-) diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex index 28314c44..ccb89c75 100644 --- a/lib/bumblebee.ex +++ b/lib/bumblebee.ex @@ -336,10 +336,22 @@ defmodule Bumblebee do * `:subdir` - the directory within the repository where the files are located - * `{:local, directory}` - the directory containing model files + * `:recursive` - whether to list files recursively. Defaults to `true` + + * `{:local, directory}` or `{:local, directory, options}` - the directory containing model files. + The supported options are: + + * `:subdir` - the directory within the directory where the + files are located + + * `:recursive` - whether to list files recursively. Defaults to `true` """ - @type repository :: {:hf, String.t()} | {:hf, String.t(), keyword()} | {:local, Path.t()} + @type repository :: + {:hf, String.t()} + | {:hf, String.t(), keyword()} + | {:local, Path.t()} + | {:local, Path.t(), keyword()} @typedoc """ A model together with its state and metadata. @@ -1273,15 +1285,28 @@ defmodule Bumblebee do {:error, "could not infer featurizer type from the configuration"} end - defp get_repo_files({:local, dir}) do + defp get_repo_files({:local, dir, opts}) do + dir = + if subdir = opts[:subdir] do + Path.join(dir, subdir) + else + dir + end + case File.ls(dir) do {:ok, filenames} -> + paths = + if Keyword.get(opts, :recursive, true) do + Path.wildcard(Path.join(dir, "**")) + else + Enum.map(filenames, &Path.join(dir, &1)) + end + repo_files = - for filename <- filenames, - path = Path.join(dir, filename), + for path <- paths, File.regular?(path), into: %{}, - do: {filename, nil} + do: {Path.relative_to(path, dir), nil} {:ok, repo_files} @@ -1290,9 +1315,22 @@ defmodule Bumblebee do end end + defp get_repo_files({:local, dir}) when is_binary(dir) do + get_repo_files({:local, dir, []}) + end + defp get_repo_files({:hf, repository_id, opts}) do subdir = opts[:subdir] - url = HuggingFace.Hub.file_listing_url(repository_id, subdir, opts[:revision]) + recursive = Keyword.get(opts, :recursive, true) + + url = + HuggingFace.Hub.file_listing_url( + repository_id, + subdir, + opts[:revision], + recursive: recursive + ) + cache_scope = repository_id_to_cache_scope(repository_id) result = @@ -1323,8 +1361,13 @@ defmodule Bumblebee do end end - defp download({:local, dir}, filename, _etag) do - path = Path.join(dir, filename) + defp download({:local, dir, opts}, filename, _etag) do + path = + if subdir = opts[:subdir] do + Path.join([dir, subdir, filename]) + else + Path.join(dir, filename) + end if File.exists?(path) do {:ok, path} @@ -1333,6 +1376,10 @@ defmodule Bumblebee do end end + defp download({:local, dir}, filename, etag) when is_binary(dir) do + download({:local, dir, []}, filename, etag) + end + defp download({:hf, repository_id, opts}, filename, etag) do filename = if subdir = opts[:subdir] do @@ -1362,18 +1409,32 @@ defmodule Bumblebee do end defp normalize_repository!({:hf, repository_id, opts}) when is_binary(repository_id) do - opts = Keyword.validate!(opts, [:revision, :cache_dir, :offline, :auth_token, :subdir]) + opts = + Keyword.validate!(opts, [ + :revision, + :cache_dir, + :offline, + :auth_token, + :subdir, + recursive: true + ]) + {:hf, repository_id, opts} end defp normalize_repository!({:local, dir}) when is_binary(dir) do - {:local, dir} + {:local, dir, []} + end + + defp normalize_repository!({:local, dir, opts}) when is_binary(dir) and is_list(opts) do + opts = Keyword.validate!(opts, [:subdir, recursive: true]) + {:local, dir, opts} end defp normalize_repository!(other) do raise ArgumentError, - "expected repository to be either {:hf, repository_id}, {:hf, repository_id, options}" <> - " or {:local, directory}, got: #{inspect(other)}" + "expected repository to be either {:hf, repository_id}, {:hf, repository_id, options}," <> + " {:local, directory}, or {:local, directory, options}, got: #{inspect(other)}" end @doc """ diff --git a/lib/bumblebee/huggingface/hub.ex b/lib/bumblebee/huggingface/hub.ex index 6a730301..ba29dcc7 100644 --- a/lib/bumblebee/huggingface/hub.ex +++ b/lib/bumblebee/huggingface/hub.ex @@ -16,12 +16,18 @@ defmodule Bumblebee.HuggingFace.Hub do @doc """ Returns a URL to list the contents of a Hugging Face repository. + + ## Options + + * `:recursive` - whether to list repository files recursively. Defaults to `true` + """ - @spec file_listing_url(String.t(), String.t() | nil, String.t() | nil) :: String.t() - def file_listing_url(repository_id, subdir, revision) do + @spec file_listing_url(String.t(), String.t() | nil, String.t() | nil, keyword()) :: String.t() + def file_listing_url(repository_id, subdir, revision, opts \\ []) do revision = revision || "main" path = if(subdir, do: "/" <> subdir) - @huggingface_endpoint <> "/api/models/#{repository_id}/tree/#{revision}#{path}" + query = if Keyword.get(opts, :recursive, true), do: "?recursive=true", else: "" + @huggingface_endpoint <> "/api/models/#{repository_id}/tree/#{revision}#{path}#{query}" end @doc """ diff --git a/test/bumblebee/huggingface/hub_test.exs b/test/bumblebee/huggingface/hub_test.exs index 14f8a6fa..59a07b39 100644 --- a/test/bumblebee/huggingface/hub_test.exs +++ b/test/bumblebee/huggingface/hub_test.exs @@ -263,6 +263,24 @@ defmodule Bumblebee.HuggingFace.HubTest do end end + describe "file_listing_url/4" do + test "constructs recursive tree listing URL by default" do + assert Hub.file_listing_url("foo/bar", nil, nil) == + "https://huggingface.co/api/models/foo/bar/tree/main?recursive=true" + + assert Hub.file_listing_url("foo/bar", "sub/dir", "v1.0") == + "https://huggingface.co/api/models/foo/bar/tree/v1.0/sub/dir?recursive=true" + end + + test "omits recursive query param when recursive: false" do + assert Hub.file_listing_url("foo/bar", nil, nil, recursive: false) == + "https://huggingface.co/api/models/foo/bar/tree/main" + + assert Hub.file_listing_url("foo/bar", "sub/dir", "v1.0", recursive: false) == + "https://huggingface.co/api/models/foo/bar/tree/v1.0/sub/dir" + end + end + defp url(port), do: "http://localhost:#{port}" defp serve_with_etag(conn, etag, body) do diff --git a/test/bumblebee_test.exs b/test/bumblebee_test.exs index 2df4a4b3..fee4bcf3 100644 --- a/test/bumblebee_test.exs +++ b/test/bumblebee_test.exs @@ -109,5 +109,24 @@ defmodule BumblebeeTest do assert Enum.sort(Map.keys(params.data)) == Enum.sort(Map.keys(default_params.data)) end + + @tag :tmp_dir + test "loads spec from local directory containing subdirectories", %{tmp_dir: tmp_dir} do + File.mkdir_p!(Path.join(tmp_dir, "sub/dir")) + File.write!(Path.join(tmp_dir, "config.json"), ~s/{"architectures": ["BertModel"]}/) + File.write!(Path.join(tmp_dir, "sub/dir/extra.json"), ~s/{}/) + + assert {:ok, %Bumblebee.Text.Bert{}} = Bumblebee.load_spec({:local, tmp_dir}) + end + + @tag :tmp_dir + test "loads spec from local directory with :subdir option", %{tmp_dir: tmp_dir} do + sub = Path.join(tmp_dir, "sub_model") + File.mkdir_p!(sub) + File.write!(Path.join(sub, "config.json"), ~s/{"architectures": ["BertModel"]}/) + + assert {:ok, %Bumblebee.Text.Bert{}} = + Bumblebee.load_spec({:local, tmp_dir, subdir: "sub_model"}) + end end end From 0db5bd8cc5566556cfa36ff49b179084cee66211 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Thu, 1 Oct 2026 20:33:30 +0200 Subject: [PATCH 2/3] Limit local recursive file discovery to max depth of 2 by default --- lib/bumblebee.ex | 16 ++++++++++++++-- test/bumblebee_test.exs | 14 ++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex index ccb89c75..5653d700 100644 --- a/lib/bumblebee.ex +++ b/lib/bumblebee.ex @@ -346,6 +346,9 @@ defmodule Bumblebee do * `:recursive` - whether to list files recursively. Defaults to `true` + * `:max_depth` - the maximum directory depth to search when `:recursive` is `true`. + Defaults to `2` + """ @type repository :: {:hf, String.t()} @@ -1297,7 +1300,8 @@ defmodule Bumblebee do {:ok, filenames} -> paths = if Keyword.get(opts, :recursive, true) do - Path.wildcard(Path.join(dir, "**")) + max_depth = Keyword.get(opts, :max_depth, 2) + list_files_with_depth(dir, max_depth) else Enum.map(filenames, &Path.join(dir, &1)) end @@ -1361,6 +1365,14 @@ defmodule Bumblebee do end end + defp list_files_with_depth(dir, max_depth) when is_integer(max_depth) and max_depth >= 1 do + 1..max_depth + |> Enum.flat_map(fn depth -> + pattern = List.duplicate("*", depth) |> Path.join() + Path.wildcard(Path.join(dir, pattern)) + end) + end + defp download({:local, dir, opts}, filename, _etag) do path = if subdir = opts[:subdir] do @@ -1427,7 +1439,7 @@ defmodule Bumblebee do end defp normalize_repository!({:local, dir, opts}) when is_binary(dir) and is_list(opts) do - opts = Keyword.validate!(opts, [:subdir, recursive: true]) + opts = Keyword.validate!(opts, [:subdir, recursive: true, max_depth: 2]) {:local, dir, opts} end diff --git a/test/bumblebee_test.exs b/test/bumblebee_test.exs index fee4bcf3..d044fd58 100644 --- a/test/bumblebee_test.exs +++ b/test/bumblebee_test.exs @@ -128,5 +128,19 @@ defmodule BumblebeeTest do assert {:ok, %Bumblebee.Text.Bert{}} = Bumblebee.load_spec({:local, tmp_dir, subdir: "sub_model"}) end + + @tag :tmp_dir + test "limits directory traversal to :max_depth", %{tmp_dir: tmp_dir} do + deep_dir = Path.join(tmp_dir, "l1/l2/l3") + File.mkdir_p!(deep_dir) + File.write!(Path.join(deep_dir, "config.json"), ~s/{"architectures": ["BertModel"]}/) + + assert_raise ArgumentError, ~r/no config file found in the given repository/, fn -> + Bumblebee.load_spec({:local, tmp_dir, max_depth: 2}) + end + + assert {:ok, %Bumblebee.Text.Bert{}} = + Bumblebee.load_spec({:local, tmp_dir, max_depth: 3, subdir: "l1/l2/l3"}) + end end end From bc355b6fecd9f0f5deca4f2384aa526e15b1d905 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Thu, 1 Oct 2026 20:38:06 +0200 Subject: [PATCH 3/3] Fix local directory recursion depth to 2 without parameter --- lib/bumblebee.ex | 16 ++-------------- test/bumblebee_test.exs | 6 +++--- 2 files changed, 5 insertions(+), 17 deletions(-) diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex index 5653d700..5f42ba94 100644 --- a/lib/bumblebee.ex +++ b/lib/bumblebee.ex @@ -346,9 +346,6 @@ defmodule Bumblebee do * `:recursive` - whether to list files recursively. Defaults to `true` - * `:max_depth` - the maximum directory depth to search when `:recursive` is `true`. - Defaults to `2` - """ @type repository :: {:hf, String.t()} @@ -1300,8 +1297,7 @@ defmodule Bumblebee do {:ok, filenames} -> paths = if Keyword.get(opts, :recursive, true) do - max_depth = Keyword.get(opts, :max_depth, 2) - list_files_with_depth(dir, max_depth) + Path.wildcard(Path.join(dir, "*")) ++ Path.wildcard(Path.join(dir, "*/*")) else Enum.map(filenames, &Path.join(dir, &1)) end @@ -1365,14 +1361,6 @@ defmodule Bumblebee do end end - defp list_files_with_depth(dir, max_depth) when is_integer(max_depth) and max_depth >= 1 do - 1..max_depth - |> Enum.flat_map(fn depth -> - pattern = List.duplicate("*", depth) |> Path.join() - Path.wildcard(Path.join(dir, pattern)) - end) - end - defp download({:local, dir, opts}, filename, _etag) do path = if subdir = opts[:subdir] do @@ -1439,7 +1427,7 @@ defmodule Bumblebee do end defp normalize_repository!({:local, dir, opts}) when is_binary(dir) and is_list(opts) do - opts = Keyword.validate!(opts, [:subdir, recursive: true, max_depth: 2]) + opts = Keyword.validate!(opts, [:subdir, recursive: true]) {:local, dir, opts} end diff --git a/test/bumblebee_test.exs b/test/bumblebee_test.exs index d044fd58..011ea579 100644 --- a/test/bumblebee_test.exs +++ b/test/bumblebee_test.exs @@ -130,17 +130,17 @@ defmodule BumblebeeTest do end @tag :tmp_dir - test "limits directory traversal to :max_depth", %{tmp_dir: tmp_dir} do + test "limits directory traversal to depth of 2", %{tmp_dir: tmp_dir} do deep_dir = Path.join(tmp_dir, "l1/l2/l3") File.mkdir_p!(deep_dir) File.write!(Path.join(deep_dir, "config.json"), ~s/{"architectures": ["BertModel"]}/) assert_raise ArgumentError, ~r/no config file found in the given repository/, fn -> - Bumblebee.load_spec({:local, tmp_dir, max_depth: 2}) + Bumblebee.load_spec({:local, tmp_dir}) end assert {:ok, %Bumblebee.Text.Bert{}} = - Bumblebee.load_spec({:local, tmp_dir, max_depth: 3, subdir: "l1/l2/l3"}) + Bumblebee.load_spec({:local, tmp_dir, subdir: "l1/l2/l3"}) end end end