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

Code before:
    tokenizer: &Tokenizer,
    args: &Args,
) -> Result<()> {
    match model {
        ModelType::UnMasked(_) => {
            println!("embeddings");
            println!("{output}");
        }
        ModelType::Masked(_) => {
            process_masked_output(output, token_ids, tokenizer, args)?;
        }
    }

    Ok(())
}

fn process_masked_output(
    output: &Tensor,
    token_ids: &Tensor,
    tokenizer: &Tokenizer,
    args: &Args,
) -> Result<()> {
    let input_ids_vec = token_ids.to_vec2::<u32>()?;
    let mask_token_id = tokenizer
        .token_to_id("[MASK]")
        .ok_or(anyhow::anyhow!("No mask token found"))?;

    println!("\nInput: {}", args.prompt);

    for (token_idx, &token_id) in input_ids_vec[0].iter().enumerate() {
        if token_id == mask_token_id {
            println!("Predictions for [MASK] at position {}:", token_idx);

            let pos_logits = output.get(0)?.get(token_idx)?;
            let probs = candle_nn::ops::softmax(&pos_logits, 0)?;
            let (top_values, top_indices) = get_top_k(&probs, args.top_k)?;

            let values = top_values.to_vec1::<f32>()?;
            let indices = top_indices.to_vec1::<u32>()?;

            for (i, (&token_id, &prob)) in indices.iter().zip(values.iter()).enumerate() {
                let token = tokenizer.decode(&[token_id], false).map_err(E::msg)?;
                println!(
                    "  {}: {:15} (probability: {:.2}%)",
                    i + 1,
                    token,
                    prob * 100.0
                );
            }
        }
    }

Review comment (suggestion):
It doesn't matter much here but `anyhow::Context` would be a bit more idiomatic there + avoid allocating the error even on the non-error path.

Code after:

    Ok((token_ids, mask))
}

fn process_output(
    model: &ModelType,
    output: &Tensor,
    token_ids: &Tensor,
    tokenizer: &Tokenizer,
    args: &Args,
) -> Result<()> {
    match model {
        ModelType::UnMasked(_) => {
            println!("embeddings");
            println!("{output}");
        }
        ModelType::Masked(_) => {
            process_masked_output(output, token_ids, tokenizer, args)?;
        }
    }

    Ok(())
}

fn process_masked_output(
    output: &Tensor,
    token_ids: &Tensor,
    tokenizer: &Tokenizer,
    args: &Args,
) -> Result<()> {
    let input_ids_vec = token_ids.to_vec2::<u32>()?;
    let mask_token_id = tokenizer
        .token_to_id("[MASK]")
        .context("Mask token, \"[MASK]\", not found in tokenizer.")?;

    println!("\nInput: {}", args.prompt);

    for (token_idx, &token_id) in input_ids_vec[0].iter().enumerate() {
        if token_id == mask_token_id {
            println!("Predictions for [MASK] at position {}:", token_idx);

            let pos_logits = output.get(0)?.get(token_idx)?;
            let probs = candle_nn::ops::softmax(&pos_logits, 0)?;
            let (top_values, top_indices) = get_top_k(&probs, args.top_k)?;

            let values = top_values.to_vec1::<f32>()?;
            let indices = top_indices.to_vec1::<u32>()?;

            for (i, (&token_id, &prob)) in indices.iter().zip(values.iter()).enumerate() {
                let token = tokenizer.decode(&[token_id], false).map_err(E::msg)?;
                println!(