xberg-gliner 1.3.10

GLiNER inference used by xberg NER: span-mode ONNX runtime plus an optional Candle GLiNER2 backend.
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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

use std::collections::HashSet;
use std::path::Path;

use ort::session::Session;
use ort::session::builder::GraphOptimizationLevel;

use crate::{GlinerError, Result, RuntimeConfig};

/// GLiNER tensor input names expected by span-mode ONNX exports.
pub const INPUT_NAMES: [&str; 6] = [
    "input_ids",
    "attention_mask",
    "words_mask",
    "text_lengths",
    "span_idx",
    "span_mask",
];

/// GLiNER tensor output names expected by span-mode ONNX exports.
pub const OUTPUT_NAMES: [&str; 1] = ["logits"];

pub(crate) const TENSOR_INPUT_IDS: &str = "input_ids";
pub(crate) const TENSOR_ATTENTION_MASK: &str = "attention_mask";
pub(crate) const TENSOR_WORD_MASK: &str = "words_mask";
pub(crate) const TENSOR_TEXT_LENGTHS: &str = "text_lengths";
pub(crate) const TENSOR_SPAN_IDX: &str = "span_idx";
pub(crate) const TENSOR_SPAN_MASK: &str = "span_mask";
pub(crate) const TENSOR_LOGITS: &str = "logits";

pub(crate) fn build_session<P: AsRef<Path>>(model_path: P, runtime: &RuntimeConfig) -> Result<Session> {
    let session = Session::builder()?
        .with_optimization_level(GraphOptimizationLevel::All)
        .map_err(|error| GlinerError::Ort(ort::Error::from(error)))?
        .with_intra_threads(runtime.intra_threads)
        .map_err(|error| GlinerError::Ort(ort::Error::from(error)))?
        .with_inter_threads(1)
        .map_err(|error| GlinerError::Ort(ort::Error::from(error)))?
        .commit_from_file(model_path)?;
    Ok(session)
}

pub(crate) fn validate_session_schema(session: &Session) -> Result<()> {
    let inputs = session
        .inputs()
        .iter()
        .map(|input| input.name().to_string())
        .collect::<Vec<_>>();
    validate_schema_names("input", &INPUT_NAMES, &inputs)?;

    let outputs = session
        .outputs()
        .iter()
        .map(|output| output.name().to_string())
        .collect::<Vec<_>>();
    validate_schema_names("output", &OUTPUT_NAMES, &outputs)
}

pub(crate) fn validate_schema_names(kind: &'static str, expected: &[&'static str], actual: &[String]) -> Result<()> {
    let expected = expected.iter().copied().collect::<HashSet<_>>();
    let actual = actual.iter().map(String::as_str).collect::<HashSet<_>>();
    if expected == actual {
        return Ok(());
    }
    let mut actual = actual.into_iter().map(str::to_string).collect::<Vec<_>>();
    actual.sort_unstable();
    Err(GlinerError::UnexpectedModelSchema {
        kind,
        expected: expected.into_iter().collect(),
        actual,
    })
}