Fix typical sampling.

Ai00-X · Aug 11, 2024 · 0a2c49a · 0a2c49a
1 parent dc1bda2
commit 0a2c49a
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 11 deletions.
diff --git a/crates/ai00-core/src/sampler/mirostat.rs b/crates/ai00-core/src/sampler/mirostat.rs
@@ -1,4 +1,4 @@
-use super::{utils, Sampler};
+use super::{radix, Sampler};
 use derivative::Derivative;
 use itertools::Itertools;
 use salvo::oapi::ToSchema;
@@ -49,13 +49,13 @@ impl Sampler for MirostatSampler {
             .iter()
             .copied()
             .enumerate()
-            .map(|(id, x)| utils::F32WithIndex(id, x))
+            .map(|(id, x)| radix::F32WithIndex(id, x))
             .collect_vec();
         sorted.voracious_sort();
         let sorted = sorted
             .into_iter()
             .rev()
-            .scan((0, 0.0, 0.0), |(_, cum, _), utils::F32WithIndex(id, x)| {
+            .scan((0, 0.0, 0.0), |(_, cum, _), radix::F32WithIndex(id, x)| {
                 // if *cum > params.top_p {
                 //     None
                 // } else {

diff --git a/crates/ai00-core/src/sampler/mod.rs b/crates/ai00-core/src/sampler/mod.rs
@@ -2,7 +2,9 @@ pub mod bnf;
 pub mod mirostat;
 pub mod nucleus;
 pub mod typical;
-mod utils;
+
+mod radix;
+
 pub trait Sampler {
     /// Initialize the sampler state.
     fn init(&mut self, model_tokens: &[u16]);

diff --git a/crates/ai00-core/src/sampler/nucleus.rs b/crates/ai00-core/src/sampler/nucleus.rs
@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use super::{utils, Sampler};
+use super::{radix, Sampler};
 use derivative::Derivative;
 use itertools::Itertools;
 use salvo::oapi::ToSchema;
@@ -72,14 +72,14 @@ impl Sampler for NucleusSampler {
             .iter()
             .copied()
             .enumerate()
-            .map(|(id, x)| utils::F32WithIndex(id, x))
+            .map(|(id, x)| radix::F32WithIndex(id, x))
             .collect_vec();
         sorted.voracious_sort();
         let sorted = sorted
             .into_iter()
             .rev()
             .take(params.top_k)
-            .scan((0, 0.0, 0.0), |(_, cum, _), utils::F32WithIndex(id, x)| {
+            .scan((0, 0.0, 0.0), |(_, cum, _), radix::F32WithIndex(id, x)| {
                 if *cum > params.top_p {
                     None
                 } else {

diff --git a/crates/ai00-core/src/sampler/utils.rs → crates/ai00-core/src/sampler/radix.rs b/crates/ai00-core/src/sampler/utils.rs → crates/ai00-core/src/sampler/radix.rs
@@ -1,38 +1,49 @@
 use std::cmp::Ordering;
+
 use voracious_radix_sort::Radixable;
+
 #[derive(Copy, Clone, Debug)]
 pub struct F32WithIndex(pub usize, pub f32);
+
 impl PartialOrd for F32WithIndex {
     fn partial_cmp(&self, other: &F32WithIndex) -> Option<Ordering> {
         self.1.partial_cmp(&other.1)
     }
 }
+
 impl PartialEq for F32WithIndex {
     fn eq(&self, other: &Self) -> bool {
         self.1 == other.1
     }
 }
+
 impl Radixable<f32> for F32WithIndex {
     type Key = f32;
+
     #[inline]
     fn key(&self) -> Self::Key {
         self.1
     }
 }
+
 #[derive(Copy, Clone, Debug)]
 pub struct DoubleF32WithIndex(pub usize, pub f32, pub f32);
+
 impl PartialOrd for DoubleF32WithIndex {
     fn partial_cmp(&self, other: &DoubleF32WithIndex) -> Option<Ordering> {
         self.2.partial_cmp(&other.2)
     }
 }
+
 impl PartialEq for DoubleF32WithIndex {
     fn eq(&self, other: &Self) -> bool {
         self.2 == other.2
     }
 }
+
 impl Radixable<f32> for DoubleF32WithIndex {
     type Key = f32;
+
     #[inline]
     fn key(&self) -> Self::Key {
         self.2

diff --git a/crates/ai00-core/src/sampler/typical.rs b/crates/ai00-core/src/sampler/typical.rs
@@ -6,7 +6,7 @@ use salvo::oapi::ToSchema;
 use serde::{Deserialize, Serialize};
 use voracious_radix_sort::RadixSort;
 
-use super::{utils, Sampler};
+use super::{radix, Sampler};
 
 #[derive(Debug, Clone, Derivative, Serialize, Deserialize, ToSchema)]
 #[derivative(Default)]
@@ -79,13 +79,12 @@ impl Sampler for TypicalSampler {
         let entropy = probs.iter().map(|(_, x, y)| x * y).sum::<f32>();
         let mut sorted = probs
             .into_iter()
-            .map(|(id, x, y)| utils::DoubleF32WithIndex(id, x, (y - entropy).abs()))
+            .map(|(id, x, y)| radix::DoubleF32WithIndex(id, x, (y - entropy).abs()))
             .collect_vec();
         sorted.voracious_sort();
         let sorted = sorted
             .into_iter()
-            .rev()
-            .map(|utils::DoubleF32WithIndex(id, x, _)| (id, x))
+            .map(|radix::DoubleF32WithIndex(id, x, _)| (id, x))
             .take(params.top_k)
             .scan((0, 0.0, 0.0), |(_, cum, _), (id, x)| {
                 if *cum > params.tau {