bump deps

Author: wheynelau

Cargo.toml

@@ -16,7 +16,7 @@ log = "0.4.26"
 once_cell = "1.20.3"
 serde = { version = "1.0.218", features = ["derive"] }
 serde_json = "1.0.140"
-tokio = "1.43.0"
+tokio = "1.44.0"
 tokenizers = { version = "0.21.0", features = ["hf-hub", "http"]}
 regex = "1.11.1"
 rand = "0.9.0"

src/common.rs

@@ -34,50 +34,36 @@ fn init_string() -> Vec<&'static str> {
             .iter()
             .map(|token| {
                 let token = tokenizer.decode(&[*token], true).unwrap();
-                let token = serde_json::to_string(&token).unwrap().trim_matches('"').to_string();
+                let token = serde_json::to_string(&token)
+                    .unwrap()
+                    .trim_matches('"')
+                    .to_string();
                 let token: &'static str = Box::leak(token.into_boxed_str());
                 token
             })
             .collect()
     } else {
         // fall back to a simple whitespace tokenizer
         log::error!("Failed to load the tokenizer, falling back to a simple whitespace tokenizer");
-        contents.split_whitespace().map(|s|{
-            let s: &'static str = Box::leak(s.to_string().into_boxed_str());
-            s}
-        ).collect()
+        contents
+            .split_whitespace()
+            .map(|s| {
+                let s: &'static str = Box::leak(s.to_string().into_boxed_str());
+                s
+            })
+            .collect()
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     #[test]
-    fn test_init_string() {
-        // to be implemented
-        let tokenizer = tokenizers::Tokenizer::from_pretrained(
-            "NousResearch/DeepHermes-3-Llama-3-8B-Preview",
-            None,
-        )
-        .expect("Should have been able to load the tokenizer");
-        let contents = "This is a test";
-        let tokens = tokenizer
-            .encode(contents, false)
-            .unwrap()
-            .get_ids()
-            .to_vec();
-        let clean_tokens: Vec<String> = tokens
-            .iter()
-            .map(|token| tokenizer.decode(&[*token], true).unwrap())
-            .collect();
-        println!("{:?}", clean_tokens);
-    }
-    #[test]
-    fn test_init_string_template () {
+    fn test_init_string_template() {
         // sanity check that the strings are the same
         let baseline = raw_string();
         // because the tokens are escaped
-        let baseline =  serde_json::to_string(&baseline)
+        let baseline = serde_json::to_string(&baseline)
             .unwrap()
             .trim_matches('"')
             .to_string();

src/routes.rs

@@ -91,7 +91,8 @@ impl Default for Choice {
 }
 
 static TEMPLATE: Lazy<String> = Lazy::new(init_template);
-static RE: Lazy<regex::Regex> = Lazy::new(|| regex::Regex::new(r"\[INPUT\]|\[MAX_TOKENS\]").unwrap());
+static RE: Lazy<regex::Regex> =
+    Lazy::new(|| regex::Regex::new(r"\[INPUT\]|\[MAX_TOKENS\]").unwrap());
 
 fn init_template() -> String {
     let max_tokens = i32::MAX;
@@ -103,13 +104,12 @@ fn init_template() -> String {
 
 fn substitute_template(input: &str, max_tokens: usize, template: Option<&String>) -> String {
     let template = template.unwrap_or(&TEMPLATE);
-    RE.replace_all(template, |caps: &regex::Captures| {
-        match &caps[0] {
-            "[INPUT]" => input.to_string(),
-            "[MAX_TOKENS]" => max_tokens.to_string(),
-            _ => unreachable!(),
-        }
-    }).to_string()
+    RE.replace_all(template, |caps: &regex::Captures| match &caps[0] {
+        "[INPUT]" => input.to_string(),
+        "[MAX_TOKENS]" => max_tokens.to_string(),
+        _ => unreachable!(),
+    })
+    .to_string()
 }
 // Use the same endpoint to allow the streaming
 pub async fn common_completions(
@@ -140,7 +140,7 @@ async fn completions(_req: HttpRequest, payload: Request) -> Result<HttpResponse
         } else {
             &TOKENIZED_OUTPUT[..max_tokens].concat()
         };
-        substitute_template(&return_string, max_tokens, None)
+        substitute_template(return_string, max_tokens, None)
     } else {
         // Use the full output when max_tokens is not specified
         let return_string = &MAX_OUTPUT;

src/stream.rs

@@ -1,13 +1,13 @@
 use actix_web_lab::sse;
 use futures_util::stream::Stream;
 use once_cell::sync::Lazy;
+use rand::rngs::ThreadRng;
 use rand::Rng;
 use serde::{Deserialize, Serialize};
-use tokio::time::Sleep;
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
-use rand::rngs::ThreadRng;
+use tokio::time::Sleep;
 
 use crate::common::MAX_TOKENS;
 use crate::routes::Usage;
@@ -58,6 +58,8 @@ impl StreamingChunkResponse {
         }
     }
 }
+// TODO
+// this can be combined with the one in routes.rs
 #[derive(Deserialize, Serialize, Debug, Default)]
 struct Choice {
     index: i32,
@@ -114,18 +116,20 @@ fn init_template() -> String {
 impl Stream for StringsStream<'_> {
     type Item = Result<sse::Event, std::convert::Infallible>;
 
+    // high level
+    // Starts with state::Input
+    // switch to state::Start after a random delay
+    // Once it reaches the end of the strings, it will switch to state::Usage if log usage is enabled
+    // After state::Usage, it will switch to state::Done
+    // If log usage is not enabled, it will switch to state::Done
+    // Once it reaches state::Done, it will switch to state::Completed
+
+    // init a string for faster access
+    // let response = StreamingChunkResponse::from_string("[INPUT]".to_string());
+    // let output = serde_json::to_string(&response).unwrap();
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = &mut *self;
-        // high level
-        // Starts with state::Start
-        // Once it reaches the end of the strings, it will switch to state::Usage if log usage is enabled
-        // After state::Usage, it will switch to state::Done
-        // If log usage is not enabled, it will switch to state::Done
-        // Once it reaches state::Done, it will switch to state::Completed
 
-        // init a string for faster access
-        // let response = StreamingChunkResponse::from_string("[INPUT]".to_string());
-        // let output = serde_json::to_string(&response).unwrap();
         if let Some(sleep) = &mut this.sleep {
             if Pin::new(sleep).poll(cx).is_pending() {
                 return Poll::Pending;
@@ -134,19 +138,23 @@ impl Stream for StringsStream<'_> {
         }
 
         match this.state {
-            State:: Input => {
+            State::Input => {
                 // Input gives a fake TTFT
                 // that is your initial delay from the LLM processing the tokens
                 // this can typically be long
                 let rand = this.rng.random_range(500..1000);
-                this.sleep = Some(Box::pin(tokio::time::sleep(tokio::time::Duration::from_millis(rand))));
+                this.sleep = Some(Box::pin(tokio::time::sleep(
+                    tokio::time::Duration::from_millis(rand),
+                )));
                 this.state = State::Start;
                 Poll::Pending
             }
             State::Start => {
                 if this.index < this.max_tokens {
                     let rand = this.rng.random_range(50..100);
-                    this.sleep = Some(Box::pin(tokio::time::sleep(tokio::time::Duration::from_millis(rand))));
+                    this.sleep = Some(Box::pin(tokio::time::sleep(
+                        tokio::time::Duration::from_millis(rand),
+                    )));
                     let string_item = &this.strings[this.index];
                     this.index += 1;
                     // let chunk = StreamingChunkResponse::from_string(string_item);