neura_decoder/
decoder.rs

1use neura_graph::{Graph, Value};
2
3pub struct Decoder {
4    token: Value,
5    position: Value,
6    chosen: Value,
7    states: Vec<Value>,
8    end: u32,
9}
10
11impl Decoder {
12    pub fn greedy(
13        graph: &Graph,
14        token: Value,
15        position: Value,
16        logits: Value,
17        states: impl IntoIterator<Item = Value>,
18        end: u32,
19    ) -> Self {
20        for (held, what) in [(token, "token"), (position, "position")] {
21            assert_eq!(
22                graph.shape(held).dims(),
23                [1, 1, 1, 1],
24                "a {what} walks {:?} where one number walks [1, 1, 1, 1]",
25                graph.shape(held).dims(),
26            );
27        }
28        let chosen = graph.argmax(logits);
29        graph.retain(chosen);
30        Self {
31            token,
32            position,
33            chosen,
34            states: states.into_iter().collect(),
35            end,
36        }
37    }
38
39    pub fn token(&self) -> Value {
40        self.token
41    }
42
43    pub fn position(&self) -> Value {
44        self.position
45    }
46
47    pub fn chosen(&self) -> Value {
48        self.chosen
49    }
50
51    pub fn states(&self) -> &[Value] {
52        &self.states
53    }
54
55    pub fn end(&self) -> u32 {
56        self.end
57    }
58}
neura_decoder/decoder.rs

neura_decoder/
decoder.rs
