Repository: huggingface/candle
File: candle-examples/examples/distilbert/main.rs (Rust)

Code before:
    ) -> Result<(PathBuf, PathBuf, PathBuf)> {
        let repo = Repo::with_revision(model_id.to_string(), RepoType::Model, revision.to_string());
        let api = Api::new()?;
        let api = api.repo(repo);

        let config = api.get("config.json")?;
        let tokenizer = api.get("tokenizer.json")?;
        let weights = if self.use_pth {
            api.get("pytorch_model.bin")?
        } else {
            api.get("model.safetensors")?
        };

        Ok((config, tokenizer, weights))
    }

    fn load_variables(&self, weights_path: &PathBuf, device: &Device) -> Result<VarBuilder> {
        if self.use_pth {
            Ok(VarBuilder::from_pth(weights_path, DTYPE, device)?)
        } else {
            Ok(unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DTYPE, device)? })
        }
    }

    fn create_model(&self, config: &Config, vb: VarBuilder) -> Result<ModelType> {
        if self.prompt.contains("[MASK]") {
            Ok(ModelType::Masked(DistilBertForMaskedLM::load(vb, config)?))
        } else {
            Ok(ModelType::UnMasked(DistilBertModel::load(vb, config)?))
        }
    }
}

fn main() -> Result<()> {
    let args = Args::parse();
    let _guard = setup_tracing(&args);

    let (model, tokenizer) = args.build_model_and_tokenizer()?;
    let device = model.device();

    let (token_ids, mask) = prepare_inputs(&args, &tokenizer, device)?;
    let output = model.forward(&token_ids, &mask)?;

    process_output(&model, &output, &token_ids, &tokenizer, &args)?;

    Ok(())
}

fn setup_tracing(args: &Args) -> Option<impl Drop> {
    if args.tracing {
        use tracing_chrome::ChromeLayerBuilder;

Review comment (suggestion):
Matching on the string seems a bit error-prone, I would be in favor of having an explicit command line flag to trigger this behavior.

Code after:
            (None, None) => (default_model, default_revision),
        }
    }

    fn download_model_files(
        &self,
        model_id: &str,
        revision: &str,
    ) -> Result<(PathBuf, PathBuf, PathBuf)> {
        let repo = Repo::with_revision(model_id.to_string(), RepoType::Model, revision.to_string());
        let api = Api::new()?;
        let api = api.repo(repo);

        let config = api.get("config.json")?;
        let tokenizer = api.get("tokenizer.json")?;
        let weights = if self.use_pth {
            api.get("pytorch_model.bin")?
        } else {
            api.get("model.safetensors")?
        };

        Ok((config, tokenizer, weights))
    }

    fn load_variables(&self, weights_path: &PathBuf, device: &Device) -> Result<VarBuilder> {
        if self.use_pth {
            Ok(VarBuilder::from_pth(weights_path, DTYPE, device)?)
        } else {
            Ok(unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DTYPE, device)? })
        }
    }

    fn create_model(&self, config: &Config, vb: VarBuilder) -> Result<ModelType> {
        match self.model {
            Which::DistilbertForMaskedLM => {
                Ok(ModelType::Masked(DistilBertForMaskedLM::load(vb, config)?))
            }
            Which::DistilBert => Ok(ModelType::UnMasked(DistilBertModel::load(vb, config)?)),
        }
    }
}

fn main() -> Result<()> {
    let args = Args::parse();
    let _guard = setup_tracing(&args);

    let (model, tokenizer) = args.build_model_and_tokenizer()?;
    let device = model.device();

    let (token_ids, mask) = prepare_inputs(&args, &tokenizer, device)?;
    let output = model.forward(&token_ids, &mask)?;