Support recursive repository file discovery and options for local repositories

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, "*")) ++ 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 """

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 """

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

test/bumblebee_test.exs

@@ -109,5 +109,38 @@ 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
+
+    @tag :tmp_dir
+    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})
+      end
+
+      assert {:ok, %Bumblebee.Text.Bert{}} =
+               Bumblebee.load_spec({:local, tmp_dir, subdir: "l1/l2/l3"})
+    end
   end
 end