package kaun

  1. Overview
  2. Docs
Flax-inspired neural network library for OCaml

Install 

dune-project
 Dependency

github.com Readme Changelog Edit opam file Versions (3)

Authors

  1. Thibaut Mattio

Maintainers

  1. Thibaut Mattio

Sources

raven-1.0.0.alpha2.tbz
sha256=93abc49d075a1754442ccf495645bc4fdc83e4c66391ec8aca8fa15d2b4f44d2
sha512=5eb958c51f30ae46abded4c96f48d1825f79c7ce03f975f9a6237cdfed0d62c0b4a0774296694def391573d849d1f869919c49008acffca95946b818ad325f6f

doc/kaun.huggingface/Kaun_huggingface/index.html

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 -> unit -> Kaun.params download_result

load_safetensors ~model_id ~dtype () downloads and loads safetensors weights. Checks for sharded checkpoints via "model.safetensors.index.json" (and similar) and merges the shards, falling back to single-file safetensors names only when no index is present.

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 -> unit -> 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