How to use {:local, dir} when loading models in Bumblebee?

Questions / Help
1

I’m getting errors when trying to use {:local, cache_dir} in Bumblebee’s load_model function. Checking the cache directory using cache_dir(), I was able to locate the models that I previously downloaded, but the files inside are hashed. Is there any way I could use the downloaded models so I can call each model on demand and not use too much memory/storage?

Here is an example code that I use to load from the cache if the model is already downloaded:

defp maybe_load_from_cache(whisper_model) do
    {_, model} = whisper_model
    model_path = String.replace(model, "/", "--")
    cache_dir = Path.join(Bumblebee.cache_dir(), "huggingface/#{model_path}")

    case File.ls(cache_dir) do
      {:ok, _files} ->
        {:local, cache_dir}
      _ ->
        whisper_model
    end
  end
2

Hey @exfrmdlrstu! There are two distinct bits:

  1. The :local source allows you to point to a local directory that follows the usual structure of a HuggingFace Hub repository. For example, you could git clone a repository such as google-bert/bert-base-uncased and then use {:local, path_to_dir} to load from there.

  2. The internal cache, where bumblebee downloads the individual files from a remote repository. This cache is opaque and Bumblebee uses it to only download the given model/tokenizer file once. When you call Bumblebee.load_model/1 second time, it only makes a quick request to the HuggingFace repo to make sure the remote file hasn’t changed, and if it hasn’t, then simply reuses the one downloaded previously. You can actually disable all networking by setting BUMBLEBEE_OFFLINE=1 env var, or by loading from {:hf, "org/repo", offline: true} (if the model is not in the cache yet, the function raises).

Looking at your snippet specifically, I think what you are trying to do is what Bumblebee.load_model/1 already does by default, as described above. Let me know if I missed something : )

2 Likes
3

I see, so that’s what Bumblebee’s cache is for. This makes a lot of sense. Looks like I don’t need to worry using load_model on demand. Thank you for this, @jonatanklosko !

1 Like