Module: OnnxRuby::Hub

Defined in:
lib/onnx_ruby/hub.rb

Constant Summary collapse

DEFAULT_CACHE_DIR = 
File.join(Dir.home, ".cache", "onnx_ruby", "models")

Class Method Summary collapse

Class Method Details

.cached_models(cache_dir: DEFAULT_CACHE_DIR) ⇒ Array<String>

List cached models

Parameters:

  • cache_dir (String) (defaults to: DEFAULT_CACHE_DIR) —

    cache directory to search

Returns:

  • (Array<String>) —

    list of cached model paths



34
35
36
37
38
 
# File 'lib/onnx_ruby/hub.rb', line 34

def self.cached_models(cache_dir: DEFAULT_CACHE_DIR)
  return [] unless Dir.exist?(cache_dir)

  Dir.glob(File.join(cache_dir, "**", "*.onnx"))
end

.clear_cache(cache_dir: DEFAULT_CACHE_DIR) ⇒ Object

Clear the model cache

Parameters:

  • cache_dir (String) (defaults to: DEFAULT_CACHE_DIR) —

    cache directory to clear



42
43
44
 
# File 'lib/onnx_ruby/hub.rb', line 42

def self.clear_cache(cache_dir: DEFAULT_CACHE_DIR)
  FileUtils.rm_rf(cache_dir) if Dir.exist?(cache_dir)
end

.download(repo_id, filename: "model.onnx", cache_dir: DEFAULT_CACHE_DIR, revision: "main") ⇒ String

Download a model from Hugging Face Hub

Parameters:

  • repo_id (String) —

    e.g. "sentence-transformers/all-MiniLM-L6-v2"

  • filename (String) (defaults to: "model.onnx") —

    ONNX file to download (default: "model.onnx")

  • cache_dir (String) (defaults to: DEFAULT_CACHE_DIR) —

    local cache directory

Returns:

  • (String) —

    path to the downloaded model file



17
18
19
20
21
22
23
24
25
26
27
28
29
 
# File 'lib/onnx_ruby/hub.rb', line 17

def self.download(repo_id, filename: "model.onnx", cache_dir: DEFAULT_CACHE_DIR, revision: "main")
  model_dir = File.join(cache_dir, repo_id.tr("/", "--"), revision)
  model_path = File.join(model_dir, filename)

  return model_path if File.exist?(model_path)

  FileUtils.mkdir_p(model_dir)

  url = "https://huggingface.co/#{repo_id}/resolve/#{revision}/#{filename}"
  download_file(url, model_path)

  model_path
end