Add top_k in nucleus sampler.

crates/ai00-core/src/sampler/nucleus.rs

@@ -9,44 +9,22 @@ use super::Sampler;
 
 #[derive(Debug, Clone, Derivative, Serialize, Deserialize, ToSchema)]
 #[derivative(Default)]
+#[serde(default)]
 pub struct NucleusParams {
     #[derivative(Default(value = "0.5"))]
-    #[serde(default = "default_top_p")]
     pub top_p: f32,
+    #[derivative(Default(value = "128"))]
+    pub top_k: usize,
     #[derivative(Default(value = "1.0"))]
-    #[serde(default = "default_temperature")]
     pub temperature: f32,
     #[derivative(Default(value = "0.3"))]
-    #[serde(default = "default_presence_penalty")]
     pub presence_penalty: f32,
     #[derivative(Default(value = "0.3"))]
-    #[serde(default = "default_frequency_penalty")]
     pub frequency_penalty: f32,
     #[derivative(Default(value = "0.99654026"))]
-    #[serde(default = "default_penalty_decay")]
     pub penalty_decay: f32,
 }
 
-fn default_top_p() -> f32 {
-    NucleusParams::default().top_p
-}
-
-fn default_temperature() -> f32 {
-    NucleusParams::default().temperature
-}
-
-fn default_presence_penalty() -> f32 {
-    NucleusParams::default().presence_penalty
-}
-
-fn default_frequency_penalty() -> f32 {
-    NucleusParams::default().frequency_penalty
-}
-
-fn default_penalty_decay() -> f32 {
-    NucleusParams::default().penalty_decay
-}
-
 #[derive(Debug, Default, Clone)]
 pub struct NucleusState {
     pub penalties: HashMap<u16, f32>,
@@ -95,6 +73,7 @@ impl Sampler for NucleusSampler {
             .iter()
             .enumerate()
             .sorted_unstable_by(|(_, x), (_, y)| x.total_cmp(y).reverse())
+            .take(params.top_k)
             .scan((0, 0.0, 0.0), |(_, cum, _), (id, x)| {
                 if *cum > params.top_p {
                     None