muna 0.0.12

Run prediction functions in your Rust apps.
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

/*
*   Muna
*   Copyright © 2026 NatML Inc. All Rights Reserved.
*/

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use futures_core::Stream;

use crate::client::{MunaClient, Result};
use crate::types::{Acceleration, Prediction, Value};
use super::local::LocalPredictionService;
use super::remote::RemotePredictionService;

/// Make predictions.
#[derive(Clone)]
pub struct PredictionService {
    local: Arc<LocalPredictionService>,
    remote: Arc<RemotePredictionService>,
}

impl PredictionService {

    pub fn new(client: Arc<MunaClient>) -> Self {
        Self {
            local: Arc::new(LocalPredictionService::new(client.clone())),
            remote: Arc::new(RemotePredictionService::new(client)),
        }
    }

    /// Create a prediction.
    pub async fn create(
        &self,
        tag: &str,
        inputs: Option<HashMap<String, Value>>,
        acceleration: Option<Acceleration>,
        client_id: Option<String>,
        configuration_id: Option<String>,
    ) -> Result<Prediction> {
        let is_local = inputs.is_none() || match acceleration.as_ref() {
            Some(Acceleration::LocalAuto | Acceleration::LocalCpu | Acceleration::LocalGpu | Acceleration::LocalNpu) => true,
            Some(Acceleration::Adaptive(s)) => s.starts_with("local_"),
            None => true,
            _ => false,
        };
        if is_local {
            self.local.create(tag, inputs, acceleration, client_id, configuration_id).await
        } else {
            self.remote.create(tag, &inputs.unwrap(), acceleration).await
        }
    }

    /// Stream a prediction.
    pub async fn stream(
        &self,
        tag: &str,
        inputs: HashMap<String, Value>,
        acceleration: Option<Acceleration>,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<Prediction>> + Send>>> {
        let is_local = match acceleration.as_ref() {
            Some(Acceleration::LocalAuto | Acceleration::LocalCpu | Acceleration::LocalGpu | Acceleration::LocalNpu) => true,
            Some(Acceleration::Adaptive(s)) => s.starts_with("local_"),
            None => true,
            _ => false,
        };
        if is_local {
            self.local.stream(tag, inputs, acceleration).await
        } else {
            self.remote.stream(tag, &inputs, acceleration).await
        }
    }

    /// Delete a predictor that is loaded in memory.
    pub async fn delete(&self, tag: &str) -> Result<bool> {
        self.local.delete(tag).await
    }
}