rten-generate 0.12.0

Utilities to simplify running auto-regressive models with RTen
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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

//! Abstraction over [`rten::Model`] for querying and executing ML models.

use std::error::Error;

use rten::{Dimension, InputOrOutput, NodeId, Output};

/// Describes the name and shape of a model input or output.
///
/// This is similar to [`rten::NodeInfo`] but the name and shape are required.
#[derive(Clone)]
pub struct NodeInfo {
    name: String,
    shape: Vec<Dimension>,
}

impl NodeInfo {
    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn shape(&self) -> &[Dimension] {
        &self.shape
    }

    pub fn from_name_shape(name: &str, shape: &[Dimension]) -> NodeInfo {
        NodeInfo {
            name: name.to_string(),
            shape: shape.to_vec(),
        }
    }
}

/// Abstraction over [`rten::Model`] used by [`Generator`](crate::Generator) to
/// query and execute a machine learning model.
///
/// This is implemented by [`rten::Model`] and the trait's methods correspond
/// to methods of the same name in that type.
pub trait Model {
    /// Get the ID of an input or output node.
    fn find_node(&self, name: &str) -> Option<NodeId>;

    /// Get the name and shape of an input or output node.
    ///
    /// Returns `None` if the node does not exist, or name or shape information
    /// is not available.
    fn node_info(&self, id: NodeId) -> Option<NodeInfo>;

    /// Return the node IDs of the model's inputs.
    fn input_ids(&self) -> &[NodeId];

    /// Run the model with the provided inputs and return the results.
    fn run(
        &self,
        inputs: Vec<(NodeId, InputOrOutput)>,
        outputs: &[NodeId],
    ) -> Result<Vec<Output>, Box<dyn Error>>;

    /// Run as much of the model as possible given the provided inputs and
    /// return the leaves of the evaluation where execution stopped.
    fn partial_run(
        &self,
        inputs: Vec<(NodeId, InputOrOutput)>,
        outputs: &[NodeId],
    ) -> Result<Vec<(NodeId, Output)>, Box<dyn Error>>;
}

impl Model for rten::Model {
    fn find_node(&self, name: &str) -> Option<NodeId> {
        self.find_node(name)
    }

    fn node_info(&self, id: NodeId) -> Option<NodeInfo> {
        self.node_info(id).and_then(|info| {
            let name = info.name()?;
            let dims = info.shape()?;

            Some(NodeInfo {
                name: name.to_string(),
                shape: dims,
            })
        })
    }

    fn input_ids(&self) -> &[NodeId] {
        self.input_ids()
    }

    fn run(
        &self,
        inputs: Vec<(NodeId, InputOrOutput)>,
        outputs: &[NodeId],
    ) -> Result<Vec<Output>, Box<dyn Error>> {
        self.run(inputs, outputs, None).map_err(|e| e.into())
    }

    fn partial_run(
        &self,
        inputs: Vec<(NodeId, InputOrOutput)>,
        outputs: &[NodeId],
    ) -> Result<Vec<(NodeId, Output)>, Box<dyn Error>> {
        self.partial_run(inputs, outputs, None)
            .map_err(|e| e.into())
    }
}