ringo-fm 0.1.0

Idiomatic Rust bindings for Apple's on-device Foundation Models framework
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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! Generation options & sampling modes.

use serde::Serialize;

use crate::error::{Error, Result};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SamplingModeKind {
    Greedy,
    Random,
}

impl SamplingModeKind {
    fn as_str(self) -> &'static str {
        match self {
            SamplingModeKind::Greedy => "greedy",
            SamplingModeKind::Random => "random",
        }
    }
}

/// Sampling strategy.
#[derive(Debug, Clone, Default)]
pub struct SamplingMode {
    pub kind: Option<SamplingModeKind>,
    pub top: Option<u32>,
    pub probability_threshold: Option<f64>,
    pub seed: Option<u64>,
}

impl SamplingMode {
    pub fn greedy() -> Self {
        Self { kind: Some(SamplingModeKind::Greedy), ..Default::default() }
    }

    /// Random sampling. Either `top` (top-k) or `probability_threshold` may be set, not both.
    pub fn random(top: Option<u32>, probability_threshold: Option<f64>, seed: Option<u64>) -> Result<Self> {
        if top.is_some() && probability_threshold.is_some() {
            return Err(Error::Native(
                "Cannot specify both 'top' and 'probability_threshold'".into(),
            ));
        }
        if let Some(p) = probability_threshold {
            if !(0.0..=1.0).contains(&p) {
                return Err(Error::Native("'probability_threshold' must be in [0.0, 1.0]".into()));
            }
        }
        Ok(Self {
            kind: Some(SamplingModeKind::Random),
            top,
            probability_threshold,
            seed,
        })
    }
}

#[derive(Debug, Clone, Default)]
pub struct GenerationOptions {
    pub sampling: Option<SamplingMode>,
    pub temperature: Option<f64>,
    pub maximum_response_tokens: Option<u32>,
}

impl GenerationOptions {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_temperature(mut self, t: f64) -> Self {
        self.temperature = Some(t);
        self
    }

    pub fn with_maximum_response_tokens(mut self, n: u32) -> Self {
        self.maximum_response_tokens = Some(n);
        self
    }

    pub fn with_sampling(mut self, s: SamplingMode) -> Self {
        self.sampling = Some(s);
        self
    }

    /// Render to the JSON form expected by the C bridge (mirrors `to_dict` in Python).
    pub(crate) fn to_json(&self) -> Result<Option<String>> {
        if self.sampling.is_none() && self.temperature.is_none() && self.maximum_response_tokens.is_none() {
            return Ok(None);
        }
        #[derive(Serialize)]
        struct Sampling<'a> {
            mode: &'a str,
            #[serde(skip_serializing_if = "Option::is_none")]
            top_k: Option<String>,
            #[serde(skip_serializing_if = "Option::is_none")]
            top_p: Option<String>,
            #[serde(skip_serializing_if = "Option::is_none")]
            seed: Option<String>,
        }
        #[derive(Serialize, Default)]
        struct Opts<'a> {
            #[serde(skip_serializing_if = "Option::is_none")]
            sampling: Option<Sampling<'a>>,
            #[serde(skip_serializing_if = "Option::is_none")]
            temperature: Option<f64>,
            #[serde(skip_serializing_if = "Option::is_none")]
            maximum_response_tokens: Option<u32>,
        }

        let sampling = self.sampling.as_ref().and_then(|s| {
            let kind = s.kind?;
            let (top_k, top_p, seed) = if kind == SamplingModeKind::Random {
                (
                    s.top.map(|v| v.to_string()),
                    s.probability_threshold.map(|v| v.to_string()),
                    s.seed.map(|v| v.to_string()),
                )
            } else {
                (None, None, None)
            };
            Some(Sampling { mode: kind.as_str(), top_k, top_p, seed })
        });
        let opts = Opts {
            sampling,
            temperature: self.temperature,
            maximum_response_tokens: self.maximum_response_tokens,
        };
        Ok(Some(serde_json::to_string(&opts)?))
    }
}