Struct grpc_vision_svc::image_captioning::model_loader::ModelLoader

pub struct ModelLoader<T: ModelLoaderApi> { /* private fields */ }

ModelLoader is a struct used to load models from the Hugging Face API.

Implementations

impl<T: ModelLoaderApi> ModelLoader<T>

pub fn new(api: T) -> Self

Creates a new instance of ModelLoader with the provided API.

Parameters

• api: An instance of a type implementing the ModelLoaderApi trait.

Returns

A new ModelLoader instance.

Example

let api = ApiBuilder::new()
    .with_token(Some("API_TOKEN".into()))
    .with_cache_dir(PathBuf::from("./cache/models"))
    .build()
    .unwrap();

let loader = ModelLoader::new(api);

pub fn load(&self, model_cfg: &ModelConfig) -> Result<Model>

Loads a model from the Hugging Face API based on the provided ModelConfig.

Parameters

• model_cfg: A reference to a ModelConfig that specifies the model to load.

Returns

A Model struct containing the paths to the model and tokenizer files.

Example

let config = ModelConfig {
    repository: "google-bert/bert-base-uncased".to_string(),
    revision: None,
    model: "model.safetensors".to_string(),
    tokenizer: "tokenizer.json".to_string(),
};
let api = ApiBuilder::new()
    .with_token(Some("API_TOKEN".into()))
    .with_cache_dir(PathBuf::from("./cache/models"))
    .build()
    .unwrap();

let loader = ModelLoader::new(api);
let model = loader.load(&config).unwrap();

pub fn load_from_toml<P: AsRef<Path>>(&self, path: P) -> Result<Models>

Loads models specified in a TOML configuration file. The TOML file should contain one or more [[model]] tables, each specifying a model to load. Each table is deserialized into a ModelConfig instance and used to load a model from the Hugging Face API. The loaded models are returned in a Models instance (a HashMap).

Parameters

• path: A reference to a Path pointing to the TOML configuration file.

Returns

A HashMap where the keys are the repository names and the values are the corresponding Model structs.

Example

let api = Api::new().unwrap();
let loader = ModelLoader::new(api);
let models = loader.load_from_toml("models.toml").unwrap();

assert!(models.get("Salesforce/blip-image-captioning-large").is_some());
assert!(models.get("microsoft/kosmos-2-patch14-224").is_some());

Example TOML config file

[[model]]
repository = "Salesforce/blip-image-captioning-large"
revision = "refs/pr/18" # Optional
model = "model.safetensors"
tokenizer = "tokenizer.json"

[[model]]
repository = "microsoft/kosmos-2-patch14-224"
model = "model.safetensors"
tokenizer = "tokenizer.json"