Downloading a model to a local Directory

[kj3moraes]

Hello there, I recently started using rust-bert and I wanted to know what the right way was to

• check if a local copy exists in my cache directory (if it does read from there)
• if it does not, download from a remote resource and save it to a custom directory.

Currently I am doing something like this

 let mut model_cache_path = PathBuf::from(MODELS_CACHE_DIR);
        model_cache_path.push("albert");

        let model = if model_cache_path.exists() {
            println!("Creating from local cached model");
            SentenceEmbeddingsBuilder::local(model_cache_path)
                .create_model()
                .expect("Failed to create the AlBERT model")
        } else {
            fs::create_dir_all(&model_cache_path).unwrap();

            // Define the url where the model will be obtained and the directory
            // that it will be cached to
            let remote_resource = RemoteResource::new(
                SentenceEmbeddingsConfigResources::PARAPHRASE_ALBERT_SMALL_V2.1,
                model_cache_path.to_str().unwrap(),
            );

            // Get the model from a remote resource
            let model = SentenceEmbeddingsBuilder::remote(
                SentenceEmbeddingsModelType::ParaphraseAlbertSmallV2,
            )
            .create_model()
            .expect("Failed to create the AlBERT model");

             // Save the model to the cache here 
            // ??

            model
        };

This code would not save to my MODELS_CACHE_DIR. Using RemoteResources to specify modules_config, transformer_config, etc. will not download these to that specific directory, it just declares these paths.

[blmarket]

Model is lazy-downloaded - you need to use model at least once to make sure you download the model. Minimal example code to download the model:

    #[test]
    fn test_load_using_bert() -> anyhow::Result<()> {
        let model = SentenceEmbeddingsModel::new(SentenceEmbeddingsConfig::from(SentenceEmbeddingsModelType::ParaphraseAlbertSmallV2))?;
        dbg!(model.encode(&["This is a test sentence."])?);
        Ok(())
    }

[kj3moraes]

Fair enough. Is there a way to say at model initialization that when the model is downloaded I want it downloaded to my_cache_dir instead of ~/.cache/.rustbert ?

[guillaume-be]

Models get downloaded to the RUSTBERT_CACHE environment variable if set

[kj3moraes]

Ok, thanks that solves it.