diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex index 28314c44..5f42ba94 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, "*")) ++ 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..011ea579 100644 --- a/test/bumblebee_test.exs +++ b/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