model_download.jl

using Pkg.Artifacts                                                                                                                                                               

const ARTIFACTS_TOML = joinpath(@__DIR__, "Artifacts.toml")

const CLOUDFRONT_DISTRIB_PREFIX = "https://cdn.huggingface.co"
const S3_BUCKET_PREFIX = "https://s3.amazonaws.com/models.huggingface.co/bert"
const DEFAULT_CONFIG_NAME = "config.json"
const DEFAULT_WEIGHT_NAME = "pytorch_model.bin"

function register_model_artifact(name; config=DEFAULT_CONFIG_NAME, weight=DEFAULT_WEIGHT_NAME, download_prefix=CLOUDFRONT_DISTRIB_PREFIX)
  global ARTIFACTS_TOML                                                                                                                                                                       
  model_hash = artifact_hash(name, ARTIFACTS_TOML)
  islegacy = occursin('/', name)

  if isnothing(model_hash) || !artifact_exist(model_hash)
    model_hash = create_artifact() do artifact_dir
      if islegacy
        download(join((download_prefix, config), '-'), joinpath(artifact_dir, config))
        download(join((download_prefix, weight), '-'), joinpath(artifact_dir, weight))
      else
        download(joinpath(download_prefix, config), joinpath(artifact_dir, config))
        download(joinpath(download_prefix, weight), joinpath(artifact_dir, weight))
      end
    end

    bind_artifact!(ARTIFACTS_TOML, name, model_hash)
  end

  model_hash
end