package kaun

  1. Overview
  2. Docs

Module Kaun_huggingfaceSource

Kaun_huggingface: HuggingFace model hub integration for Kaun.

Provides seamless loading of pretrained models from HuggingFace, with automatic downloading, caching, and type-safe model definitions.

Types

Sourcetype model_id = string

HuggingFace model identifier (e.g., "gpt2", "bert-base-uncased")

Sourcetype revision =
  1. | Latest
  2. | Tag of string
  3. | Commit of string
    (*

    Model revision specification

    *)
Sourcetype cache_dir = string

Local cache directory for downloaded models

Sourcetype download_progress = {
  1. downloaded_bytes : int;
  2. total_bytes : int option;
  3. rate : float;
}
Sourcetype 'a download_result =
  1. | Cached of 'a
  2. | Downloaded of 'a * download_progress

Configuration

Sourcemodule Config : sig ... end

Model Registry

Sourcemodule Registry : sig ... end

Registry of known model architectures with their loading functions

Core Loading Functions

Sourceval download_file : ?config:Config.t -> ?revision:revision -> model_id:model_id -> filename:string -> unit -> string download_result

download_file ~model_id ~filename () downloads a single file from HuggingFace. Returns the local path to the file (either cached or newly downloaded).

Sourceval load_safetensors : ?config:Config.t -> ?revision:revision -> model_id:model_id -> dtype:(float, 'a) Rune.dtype -> unit -> 'a Kaun.params download_result

load_safetensors ~model_id ~dtype () downloads and loads safetensors weights. Automatically tries common filenames like "model.safetensors".

Sourceval load_config : ?config:Config.t -> ?revision:revision -> model_id:model_id -> unit -> Yojson.Safe.t download_result

load_config ~model_id () downloads and parses the model's config.json

High-level Model Loading

Sourceval from_pretrained : ?config:Config.t -> ?revision:revision -> model_id:model_id -> dtype:(float, 'a) Rune.dtype -> unit -> 'a Kaun.params

from_pretrained ~model_id ~dtype () loads a complete model.

This is the main entry point for loading models. It: 1. Downloads the model configuration 2. Downloads the model weights 3. Loads and returns the parameter tree

  • raises Failure

    if model architecture is unknown or download fails

Example:

  let gpt2_params =
    Kaun_huggingface.from_pretrained ~model_id:"gpt2" ~dtype:Rune.Float32 ()

Utilities

Sourceval list_cached_models : ?config:Config.t -> unit -> model_id list

List all models in the local cache

Sourceval clear_cache : ?config:Config.t -> ?model_id:model_id -> unit -> unit

Clear the cache for a specific model or all models

Sourceval get_model_info : model_id -> (Yojson.Safe.t, string) result

Fetch model card/info from HuggingFace API