Better way to do SamplerParams.

src/api/oai/mod.rs

@@ -19,21 +19,34 @@ use crate::sampler::{
     Sampler,
 };
 
-#[derive(Debug, Default, Clone, Deserialize)]
-pub struct SamplerParams {
-    #[serde(flatten)]
-    nucleus: Option<NucleusParams>,
-    #[serde(flatten)]
-    mirostat: Option<MirostatParams>,
+#[derive(Debug, Clone, Deserialize)]
+#[serde(untagged)]
+pub enum SamplerParams {
+    Nucleus {
+        #[serde(flatten)]
+        params: NucleusParams,
+    },
+    Mirostat {
+        #[serde(flatten)]
+        params: MirostatParams,
+    },
+}
+
+impl Default for SamplerParams {
+    fn default() -> Self {
+        Self::Nucleus {
+            params: Default::default(),
+        }
+    }
 }
 
 impl From<SamplerParams> for Arc<RwLock<dyn Sampler + Send + Sync>> {
     fn from(value: SamplerParams) -> Self {
-        let SamplerParams { nucleus, mirostat } = value;
-        match (nucleus, mirostat) {
-            (None, None) => Arc::new(RwLock::new(NucleusSampler::new(Default::default()))),
-            (None, Some(params)) => Arc::new(RwLock::new(MirostatSampler::new(params))),
-            (Some(params), _) => Arc::new(RwLock::new(NucleusSampler::new(params))),
+        match value {
+            SamplerParams::Nucleus { params } => Arc::new(RwLock::new(NucleusSampler::new(params))),
+            SamplerParams::Mirostat { params } => {
+                Arc::new(RwLock::new(MirostatSampler::new(params)))
+            }
         }
     }
 }

src/sampler/mirostat.rs

@@ -6,16 +6,20 @@ use super::Sampler;
 
 #[derive(Debug, Clone, Derivative, Serialize, Deserialize)]
 #[derivative(Default)]
-#[serde(default)]
 pub struct MirostatParams {
     #[derivative(Default(value = "3.0"))]
     pub tau: f32,
     #[derivative(Default(value = "0.1"))]
     pub rate: f32,
     #[derivative(Default(value = "128"))]
+    #[serde(default = "default_threshold")]
     pub threshold: usize,
 }
 
+fn default_threshold() -> usize {
+    MirostatParams::default().threshold
+}
+
 #[derive(Debug, Clone, Default)]
 pub struct MirostatState {
     pub max_surprise: f32,