From a5f6e4ea41211fe12a295c2c1a9c9adbe691d831 Mon Sep 17 00:00:00 2001
From: cryscan <cryscan@umich.edu>
Date: Mon, 24 Jul 2023 23:28:26 +0800
Subject: [PATCH] Add logger.

---
 Cargo.toml  |  2 ++
 src/main.rs | 16 ++++++++++++----
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 4266d5f0..63e1d69a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -29,6 +29,8 @@ itertools = "0.11"
 flume = "0.10"
 futures-util = "0.3"
 anyhow = "1"
+log = "0.4"
+simple_logger = { version = "4.2", features = ["stderr"] }
 
 [profile.release]
 lto = true
diff --git a/src/main.rs b/src/main.rs
index 3f9d267f..f6992bdf 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -120,7 +120,6 @@ fn load_model(env: &Environment, path: PathBuf) -> Result<Model> {
     let file = File::open(path)?;
     let map = unsafe { Mmap::map(&file)? };
     let model = env.create_model_from_bytes(&map)?;
-    print!("{:#?}\n{:#?}\n\n", env.adapter.get_info(), model.info());
     Ok(model)
 }
 
@@ -130,9 +129,17 @@ fn model_task(
     tokenizer: PathBuf,
     receiver: Receiver<ThreadRequest>,
 ) -> Result<()> {
+    simple_logger::SimpleLogger::new()
+        .with_level(log::LevelFilter::Warn)
+        .with_module_level("ai00_server", log::LevelFilter::Trace)
+        .init()?;
+
     let tokenizer = load_tokenizer(tokenizer)?;
     let model = load_model(&env, model)?;
 
+    log::info!("{:#?}", env.adapter.get_info());
+    log::info!("{:#?}", model.info());
+
     let penalty_free_tokens = {
         let mut set = HashSet::new();
         for token in 0..u16::MAX {
@@ -185,8 +192,7 @@ fn model_task(
             RequestKind::Embedding(request) => request.into(),
         };
 
-        println!("{:#?}", sampler);
-        std::io::stdout().flush()?;
+        log::info!("{:#?}", sampler);
 
         let state = model.create_state();
         let remain = {
@@ -197,8 +203,10 @@ fn model_task(
                 .and_then(|backed| state.load(backed).ok())
                 .is_some()
             {
+                log::info!("state cache hit");
                 remain.split_off(prefix.len())
             } else {
+                log::info!("state cache miss");
                 remain
             }
         };
@@ -234,7 +242,7 @@ fn model_task(
                     .and_then(|x| String::from_utf8(x).ok())
                     .unwrap_or_default();
 
-                print!("{}", word);
+                print!("{word}");
                 std::io::stdout().flush()?;
 
                 model_text += &word;