Reorganize sampler files.

src/main.rs

@@ -22,7 +22,7 @@ use half::f16;
 use itertools::Itertools;
 use memmap2::Mmap;
 use safetensors::SafeTensors;
-use sampler::NucleusSampler;
+use sampler::nucleus::NucleusSampler;
 use serde::{Deserialize, Serialize};
 use tokio::sync::{Mutex, RwLock};
 use tower_http::{cors::CorsLayer, services::ServeDir};

src/oai/chat.rs

@@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::RwLock;
 
 use crate::{
-    sampler::{NucleusParams, NucleusSampler},
+    sampler::nucleus::{NucleusParams, NucleusSampler},
     utils::request_info,
     Array, FinishReason, GenerateRequest, ThreadRequest, ThreadState, Token, TokenCounter,
 };

src/oai/completion.rs

@@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::RwLock;
 
 use crate::{
-    sampler::{NucleusParams, NucleusSampler},
+    sampler::nucleus::{NucleusParams, NucleusSampler},
     utils::request_info,
     Array, FinishReason, GenerateRequest, ThreadRequest, ThreadState, Token, TokenCounter,
 };

src/sampler/mod.rs

@@ -0,0 +1,13 @@
+pub mod mirostat;
+pub mod nucleus;
+
+pub trait Sampler {
+    /// Initialize the sampler state.
+    fn init(&mut self, model_tokens: &[u16]);
+    /// Update the raw model output.
+    fn transform(&self, output: &mut [f32]);
+    /// Select one token from the distribution.
+    fn sample(&self, probs: &[f32]) -> u16;
+    /// Update the sampler state after a token is chosen.
+    fn update(&mut self, token: u16);
+}

src/sampler.rs → src/sampler/nucleus.rs

@@ -3,16 +3,7 @@ use std::collections::HashMap;
 use derivative::Derivative;
 use itertools::Itertools;
 
-pub trait Sampler {
-    /// Initialize the sampler state.
-    fn init(&mut self, model_tokens: &[u16]);
-    /// Update the raw model output.
-    fn transform(&self, output: &mut [f32]);
-    /// Select one token from the distribution.
-    fn sample(&self, probs: &[f32]) -> u16;
-    /// Update the sampler state after a token is chosen.
-    fn update(&mut self, token: u16);
-}
+use super::Sampler;
 
 #[derive(Debug, Clone, Derivative)]
 #[derivative(Default)]