Loading a local model using candle

[manifest]

Hey everyone!

I have fine-tuned a DistilBERT model using Transformers library. Now, I want to perform inference using Cande library.

How can I load weights stored locally using Candle?

The example isn't helpful since it shows how to work with huggingface hub only.

[manifest]

An intuitive approach produces an uninformative error

let config: Config = serde_json::from_str("path/to/checkpoint/config.json")?;
// Error: "expected value at line 1 column 1".

for the following config:

{
  "_name_or_path": "DeepPavlov/distilrubert-base-cased-conversational",
  "activation": "gelu",
  "architectures": [
    "DistilBertForSequenceClassification"
  ],
  "attention_dropout": 0.1,
  "dim": 768,
  "dropout": 0.1,
  "hidden_dim": 3072,
  "id2label": {
    "0": "NEGATIVE",
    "1": "POSITIVE"
  },
  "initializer_range": 0.02,
  "label2id": {
    "NEGATIVE": 0,
    "POSITIVE": 1
  },
  "max_position_embeddings": 512,
  "model_type": "distilbert",
  "n_heads": 12,
  "n_layers": 6,
  "output_attentions": true,
  "output_hidden_states": true,
  "output_past": true,
  "pad_token_id": 0,
  "problem_type": "single_label_classification",
  "qa_dropout": 0.1,
  "seq_classif_dropout": 0.2,
  "sinusoidal_pos_embds": false,
  "tie_weights_": true,
  "torch_dtype": "float32",
  "transformers_version": "4.36.0",
  "vocab_size": 119547
}

[manifest]

Loading with HuggingFace Hub API works.

let device = Device::Cpu;

let model_id = "DeepPavlov/distilrubert-base-cased-conversational".to_string();
let api = Api::new()?.model(model_id);

let config_filename = "path/to/checkpoint/config.json";
let config_filename = api.get(config_filename)?;
let config_filename = std::fs::read_to_string(config_filename)?;
let config: Config = serde_json::from_str(&config_filename)?;

let tokenizer_filename = "path/to/checkpoint/tokenizer.json";
let tokenizer_filename = api.get(tokenizer_filename)?;
let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(anyhow::Error::msg)?;

let parameters_filename = "path/to/checkpoint/model.safetensors";
let parameters_filename = api.get(parameters_filename)?;

let params = unsafe { VarBuilder::from_mmaped_safetensors(&[parameters_filename], DTYPE, &device)? };

let model = DistilBertModel::load(params, &config)?;

Should I always use HuggingFace Hub API even for models stored on a local disk?

[Narsil]

No you don't need it.

If you have a real path, just do

let params = unsafe { VarBuilder::from_mmaped_safetensors(&[parameters_filename], DTYPE, &device)? };

This will mmap the file for you and load it.
hf-hub is just convenient if you want to automatically download something from the hub (or load a local cached copy to get the directories right).

[manifest]

Hi @Narsil, thank you for your response.

To perform the inference, I need to initialize a model struct (DistilBertModel in my case). To do that, I need to load a config.

How can I read locally stored config.json?

An intuitive approach results in the error.

[Narsil]

Because you're trying to load the PATH as the config itself, you need to read the file first.

serde_json::from_str(std::fs::read_to_string("path/to/checkpoint/config.json")?)?;

[manifest]

True. I missed that. Thanks a lot!

[manifest]

The working example

let device = Device::Cpu;

let config_filename = "path/to/checkpoint/config.json";
let config_filename = std::fs::read_to_string(config_filename)?;
let config: Config = serde_json::from_str(&config_filename)?;

let tokenizer_filename = "path/to/checkpoint/tokenizer.json";
let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(anyhow::Error::msg)?;

let parameters_filename = "path/to/checkpoint/model.safetensors";
let params = unsafe { VarBuilder::from_mmaped_safetensors(&[parameters_filename], DTYPE, &device)? };

let model = DistilBertModel::load(params, &config)?;