embed_anything 0.7.1

Embed anything at lightning speed
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

use anyhow::{Error as E, Result};
use candle_core::{Device, IndexOp, Tensor};
use half::f16;
use ndarray::{s, Array2};
use ort::session::Session;

pub fn compute_scores(
    model: &mut Session,
    input_ids: Array2<i64>,
    attention_mask: Array2<i64>,
    position_ids: Array2<i64>,
    false_token_id: u32,
    true_token_id: u32,
) -> Result<Vec<f32>, E> {
    let input_ids_tensor = ort::value::TensorRef::from_array_view(&input_ids)?;
    let attention_mask_tensor = ort::value::TensorRef::from_array_view(&attention_mask)?;
    let position_ids_tensor = ort::value::TensorRef::from_array_view(&position_ids)?;

    let inputs = ort::inputs!["input_ids" => input_ids_tensor, "attention_mask" => attention_mask_tensor, "position_ids" => position_ids_tensor];

    let outputs = model.run(inputs)?;
    let logits = outputs["logits".to_string()]
        .try_extract_array::<f16>()?
        .to_owned()
        .into_dimensionality::<ndarray::Ix3>()?
        .map(|&x| x.to_f32())
        .into_dimensionality::<ndarray::Ix3>()?;

    let logits = logits
        .slice(s![.., input_ids.shape()[1] - 1, ..])
        .to_owned();
    let false_logits = logits.slice(s![.., false_token_id as usize]).to_owned();
    let true_logits = logits.slice(s![.., true_token_id as usize]).to_owned();

    let false_logits_tensor =
        Tensor::from_vec(false_logits.to_vec(), input_ids.shape()[0], &Device::Cpu)?;
    let true_logits_tensor =
        Tensor::from_vec(true_logits.to_vec(), input_ids.shape()[0], &Device::Cpu)?;

    let logits_tensor = Tensor::stack(&[&false_logits_tensor, &true_logits_tensor], 1)?;
    let probs = candle_nn::ops::softmax(&logits_tensor, 1)?;
    let scores = probs.i((.., 1))?;
    Ok(scores.to_vec1::<f32>()?)
}