How to extract Labels & Scores from Raw Tensor Outputs(Text Classification Model)

[bajrangCoder]

I am trying to run a PyTorch text classification model using this framework. The model successfully produces raw tensor outputs, but I'm encountering difficulty in extracting the labels and corresponding scores, similar to what other pipeline frameworks provide(in python).

My code based on candle bert example:

#[cfg(feature = "accelerate")]
extern crate accelerate_src;

use anyhow::{Error as E, Result};
use candle_core::Device;
use candle_core::Tensor;
use candle_nn::VarBuilder;
use candle_transformers::models::bert::{BertModel, Config, HiddenAct, DTYPE};
use hf_hub::{api::sync::Api, Repo, RepoType};
use tokenizers::Tokenizer;

fn build_model_and_tokenizer(
    use_pth: bool,
    approximate_gelu: bool,
) -> Result<(BertModel, Tokenizer)> {
    let model_id = "Titeiiko/OTIS-Official-Spam-Model".to_string();
    let revision = "main".to_string();
    let device = Device::Cpu;

    let repo = Repo::with_revision(model_id, RepoType::Model, revision);
    let (config_filename, tokenizer_filename, weights_filename) = {
        let api = Api::new()?;
        let api = api.repo(repo);
        let config = api.get("config.json")?;
        let tokenizer = api.get("tokenizer.json")?;
        let weights = if use_pth {
            api.get("pytorch_model.bin")?
        } else {
            api.get("model.safetensors")?
        };
        (config, tokenizer, weights)
    };

    let config = std::fs::read_to_string(config_filename)?;
    let mut config: Config = serde_json::from_str(&config)?;
    let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;

    let vb = if use_pth {
        VarBuilder::from_pth(&weights_filename, DTYPE, &device)?
    } else {
        unsafe { VarBuilder::from_mmaped_safetensors(&[weights_filename], DTYPE, &device)? }
    };

    if approximate_gelu {
        config.hidden_act = HiddenAct::GeluApproximate;
    }

    let model = BertModel::load(vb, &config)?;
    Ok((model, tokenizer))
}

fn main() -> Result<()> {
    let start = std::time::Instant::now();

    let use_pth = false;
    let approximate_gelu = false;

    let (model, mut tokenizer) = build_model_and_tokenizer(use_pth, approximate_gelu)?;
    let device = &model.device;
    let prompt = "sample spam text";
    let tokenizer = tokenizer
        .with_padding(None)
        .with_truncation(None)
        .map_err(E::msg)?;

    let tokens = tokenizer
        .encode(prompt, true)
        .map_err(E::msg)?
        .get_ids()
        .to_vec();

    let token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
    let token_type_ids = token_ids.zeros_like()?;

    println!("Loaded and encoded in {:?}", start.elapsed());

    // Perform inference.
    let start = std::time::Instant::now();
    let ys = model.forward(&token_ids, &token_type_ids)?;

    println!("Inference result: {ys}");
    println!("Took {:?}", start.elapsed());

    Ok(())
}

Any help will be appreciated!!