feat(xtask): 添加文本生成任务的初始化步骤

Signed-off-by: YdrMaster <ydrml@hotmail.com>

.cargo/config.toml

@@ -1,3 +1,4 @@
 [alias]
 xtask = "run --package xtask --release --"
 cast = "xtask cast"
+generate = "xtask generate"

model-parameters/src/lib.rs

@@ -5,8 +5,6 @@ mod save;
 #[macro_use]
 extern crate log;
 
-use std::fmt;
-
 use common::utok;
 
 pub use data_type::DataType;
@@ -63,7 +61,6 @@ pub trait Llama2 {
 }
 
 pub use memory::{Memory, SafeTensorError};
-use serde::Serialize;
 
 #[derive(serde::Serialize, serde::Deserialize, Debug)]
 struct ConfigJson {
@@ -76,20 +73,11 @@ struct ConfigJson {
     pub num_hidden_layers: usize,
     pub num_key_value_heads: usize,
     pub vocab_size: usize,
-    #[serde(serialize_with = "serialize_float")]
     pub rms_norm_eps: f32,
-    #[serde(serialize_with = "serialize_float")]
     pub rope_theta: f32,
     pub torch_dtype: DataType,
 }
 
-fn serialize_float<S>(val: &impl fmt::LowerExp, s: S) -> Result<S::Ok, S::Error>
-where
-    S: serde::Serializer,
-{
-    format!("{val:e}").serialize(s)
-}
-
 struct LayerParamsOffset {
     input_layernorm: usize,
     self_attn_q_proj: usize,

model-parameters/src/memory/safe_tensors.rs

@@ -56,7 +56,7 @@ impl Memory<Mmap> {
             let path = name.split('.').collect::<Vec<_>>();
             let offset = header_offset + tensor.data_offsets.0;
 
-            info!(target: "import safetensors", "detect {offset:#010x} -> \"{name}\"");
+            debug!(target: "import safetensors", "detect {offset:#010x} -> \"{name}\"");
             match path.as_slice() {
                 ["model", "embed_tokens", "weight"] => {
                     assert_eq!(&tensor.shape, &[config.vocab_size, d]);

transformer-cpu/src/lib.rs

@@ -3,6 +3,8 @@ mod cache;
 use cache::LayerCache;
 use model_parameters::{DataType, Llama2, Memory};
 
+pub extern crate model_parameters;
+
 pub struct Transformer {
     model: Box<dyn Llama2>,
     cache: Vec<LayerCache>,

xtask/Cargo.toml

@@ -7,5 +7,8 @@ authors = ["YdrMaster <ydrml@hotmail.com>"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-model-parameters = { path = "../model-parameters" }
+tokenizer = { path = "../tokenizer" }
+transformer-cpu = { path = "../transformer-cpu" }
+log = "0.4"
+env_logger = "0.11"
 clap = { version = "4.5", features = ["derive"] }

xtask/src/cast.rs

@@ -0,0 +1,59 @@
+use std::{fs, path::PathBuf, time::Instant};
+use transformer_cpu::model_parameters::{save, DataType, Llama2, Memory};
+
+#[derive(Args, Default)]
+pub(crate) struct CastArgs {
+    /// Original model directory.
+    #[clap(short, long)]
+    model: String,
+    /// Target model directory.
+    #[clap(short, long)]
+    target: Option<String>,
+    /// Target model type.
+    #[clap(short, long)]
+    dt: Option<String>,
+}
+
+impl CastArgs {
+    pub fn invode(self) {
+        let ty = match self.dt.as_deref() {
+            Some("f32") | Some("float") | Some("float32") | None => DataType::F32,
+            Some("f16") | Some("half") | Some("float16") => DataType::F16,
+            Some("bf16") | Some("bfloat16") => DataType::BF16,
+            Some(ty) => panic!("Unknown data type: \"{ty}\""),
+        };
+        let model_dir = PathBuf::from(self.model);
+
+        let time = Instant::now();
+        let model = Memory::load_safetensors(&model_dir).unwrap();
+        println!("load model ... {:?}", time.elapsed());
+
+        if model.data_type() == ty {
+            println!("Model already has target data type");
+            return;
+        }
+
+        let target = self.target.map(PathBuf::from).unwrap_or_else(|| {
+            model_dir.parent().unwrap().join(format!(
+                "{}_{ty:?}",
+                model_dir.file_name().unwrap().to_str().unwrap()
+            ))
+        });
+        fs::create_dir_all(&target).unwrap();
+
+        let time = Instant::now();
+        let model = Memory::cast(&model, ty);
+        println!("cast data type ... {:?}", time.elapsed());
+
+        let time = Instant::now();
+        save(&model, &target).unwrap();
+        println!("save model ... {:?}", time.elapsed());
+
+        let tokenizer = model_dir.join("tokenizer.model");
+        if tokenizer.is_file() {
+            let time = Instant::now();
+            fs::copy(&tokenizer, target.join("tokenizer.model")).unwrap();
+            println!("copy tokenizer ... {:?}", time.elapsed());
+        }
+    }
+}

xtask/src/generate.rs

@@ -0,0 +1,35 @@
+use std::{path::PathBuf, time::Instant};
+use tokenizer::{Tokenizer, BPE};
+use transformer_cpu::{model_parameters::Memory, Transformer};
+
+#[derive(Args, Default)]
+pub(crate) struct GenerateArgs {
+    /// Model directory.
+    #[clap(short, long)]
+    model: String,
+    /// Prompt.
+    #[clap(short, long)]
+    prompt: String,
+}
+
+impl GenerateArgs {
+    pub fn invoke(self) {
+        let model_dir = PathBuf::from(self.model);
+
+        let time = Instant::now();
+        let model = Box::new(Memory::load_safetensors(&model_dir).unwrap());
+        info!("load model ... {:?}", time.elapsed());
+
+        let time = Instant::now();
+        let _transformer = Transformer::new(model, 1);
+        info!("build transformer ... {:?}", time.elapsed());
+
+        let time = Instant::now();
+        let tokenizer = BPE::from_model_file(model_dir.join("tokenizer.model")).unwrap();
+        info!("build tokenizer ... {:?}", time.elapsed());
+
+        let time = Instant::now();
+        let _prompt_tokens = tokenizer.encode(self.prompt.trim());
+        info!("encode prompt ... {:?}", time.elapsed());
+    }
+}

xtask/src/main.rs

@@ -1,14 +1,21 @@
+mod cast;
+mod generate;
+
 use clap::Parser;
-use model_parameters::{save, DataType, Llama2, Memory};
-use std::{fs, path::PathBuf, time::Instant};
 
 #[macro_use]
 extern crate clap;
+#[macro_use]
+extern crate log;
 
 fn main() {
+    // set env for POWERSHELL: `$env:RUST_LOG="INFO";`
+    env_logger::init();
+
     use Commands::*;
     match Cli::parse().command {
-        Cast(args) => args.cast(),
+        Cast(cast) => cast.invode(),
+        Generate(generate) => generate.invoke(),
     }
 }
 
@@ -23,52 +30,7 @@ struct Cli {
 #[derive(Subcommand)]
 enum Commands {
     /// Cast model
-    Cast(CastArgs),
-}
-
-#[derive(Args, Default)]
-struct CastArgs {
-    /// Original model directory.
-    #[clap(short, long)]
-    model: String,
-    /// Target model directory.
-    #[clap(short, long)]
-    target: Option<String>,
-    /// Target model type.
-    #[clap(short, long)]
-    dt: Option<String>,
-}
-
-impl CastArgs {
-    fn cast(self) {
-        let ty = match self.dt.as_ref().map(String::as_str) {
-            Some("f32") | Some("float") | Some("float32") | None => DataType::F32,
-            Some("f16") | Some("half") | Some("float16") => DataType::F16,
-            Some("bf16") | Some("bfloat16") => DataType::BF16,
-            Some(ty) => panic!("Unknown data type: \"{ty}\""),
-        };
-        let model_dir = PathBuf::from(self.model);
-        let time = Instant::now();
-        let model = Memory::load_safetensors(&model_dir).unwrap();
-        println!("load model ... {:?}", time.elapsed());
-        if model.data_type() == ty {
-            println!("Model already has target data type");
-            return;
-        }
-
-        let target = self.target.map(PathBuf::from).unwrap_or_else(|| {
-            model_dir.parent().unwrap().join(format!(
-                "{}_{ty:?}",
-                model_dir.file_name().unwrap().to_str().unwrap()
-            ))
-        });
-        fs::create_dir_all(&target).unwrap();
-        let t0 = Instant::now();
-        let model = Memory::cast(&model, ty);
-        let t1 = Instant::now();
-        println!("cast data type ... {:?}", t1 - t0);
-        save(&model, target).unwrap();
-        let t2 = Instant::now();
-        println!("save model ... {:?}", t2 - t1);
-    }
+    Cast(cast::CastArgs),
+    /// Generate following text
+    Generate(generate::GenerateArgs),
 }