-
Notifications
You must be signed in to change notification settings - Fork 27
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: YdrMaster <ydrml@hotmail.com>
Showing
10 changed files
with
65 additions
and
86 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,80 +1,58 @@ | ||
| use super::ServiceArgs; | ||
| use crate::common::{argmax, logger_init, tokenizer}; | ||
| use crate::common::{argmax, tokenizer}; | ||
| use common::upos; | ||
| use std::{collections::HashMap, path::Path, time::Instant}; | ||
| use tokenizer::Tokenizer; | ||
| use transformer_cpu::{model_parameters::Memory, LayerCache, Transformer}; | ||
|
|
||
| pub(super) struct CpuService { | ||
| transformer: Transformer, | ||
| sessions: HashMap<usize, SessionContext>, | ||
| tokenizer: Box<dyn Tokenizer>, | ||
| } | ||
|
|
||
| struct SessionContext { | ||
| pos: upos, | ||
| kv_cache: Vec<LayerCache>, | ||
| } | ||
|
|
||
| impl From<ServiceArgs> for CpuService { | ||
| fn from(args: ServiceArgs) -> Self { | ||
| logger_init(args.log); | ||
|
|
||
| let model_dir = Path::new(&args.model); | ||
| pub(super) fn run(args: ServiceArgs) { | ||
| let model_dir = Path::new(&args.model); | ||
|
|
||
| let time = Instant::now(); | ||
| let tokenizer = tokenizer(args.tokenizer, &model_dir); | ||
| info!("build tokenizer ... {:?}", time.elapsed()); | ||
| let time = Instant::now(); | ||
| let tokenizer = tokenizer(args.tokenizer, &model_dir); | ||
| info!("build tokenizer ... {:?}", time.elapsed()); | ||
|
|
||
| let time = Instant::now(); | ||
| let model = Box::new(Memory::load_safetensors_from_dir(model_dir).unwrap()); | ||
| info!("load model ... {:?}", time.elapsed()); | ||
| let time = Instant::now(); | ||
| let model = Box::new(Memory::load_safetensors_from_dir(model_dir).unwrap()); | ||
| info!("load model ... {:?}", time.elapsed()); | ||
|
|
||
| let time = Instant::now(); | ||
| let transformer = Transformer::new(model); | ||
| info!("build transformer ... {:?}", time.elapsed()); | ||
| let time = Instant::now(); | ||
| let mut transformer = Transformer::new(model); | ||
| info!("build transformer ... {:?}", time.elapsed()); | ||
|
|
||
| Self { | ||
| transformer, | ||
| sessions: HashMap::new(), | ||
| tokenizer, | ||
| } | ||
| struct SessionContext { | ||
| pos: upos, | ||
| kv_cache: Vec<LayerCache>, | ||
| } | ||
| } | ||
|
|
||
| impl CpuService { | ||
| pub fn run(mut self) { | ||
| loop { | ||
| let id = 0; | ||
| let prompt = "The quick brown fox jumps over the lazy dog"; | ||
| let mut sessions = HashMap::<usize, SessionContext>::new(); | ||
|
|
||
| let session = self.sessions.entry(id).or_insert_with(|| SessionContext { | ||
| pos: 0, | ||
| kv_cache: self.transformer.new_cache(), | ||
| }); | ||
| loop { | ||
| let id = 0; | ||
| let prompt = "The quick brown fox jumps over the lazy dog"; | ||
|
|
||
| let prompt_tokens = self.tokenizer.encode(&prompt.trim()); | ||
| let (last, tokens) = prompt_tokens.split_last().expect("prompt is empty"); | ||
| if !tokens.is_empty() { | ||
| self.transformer | ||
| .update(tokens, &mut session.kv_cache, session.pos as _); | ||
| session.pos += tokens.len() as upos; | ||
| } | ||
| let session = sessions.entry(id).or_insert_with(|| SessionContext { | ||
| pos: 0, | ||
| kv_cache: transformer.new_cache(), | ||
| }); | ||
|
|
||
| let prompt_tokens = tokenizer.encode(&prompt.trim()); | ||
| let (last, tokens) = prompt_tokens.split_last().expect("prompt is empty"); | ||
| if !tokens.is_empty() { | ||
| transformer.update(tokens, &mut session.kv_cache, session.pos as _); | ||
| session.pos += tokens.len() as upos; | ||
| } | ||
|
|
||
| let mut token = *last; | ||
| let max_pos = self.transformer.max_seq_len() as upos; | ||
| let mut out = String::new(); | ||
| while session.pos < max_pos { | ||
| let logits = | ||
| self.transformer | ||
| .forward(token, &mut session.kv_cache, session.pos as _); | ||
| let next = argmax(logits); | ||
| let mut token = *last; | ||
| let max_pos = transformer.max_seq_len() as upos; | ||
| let mut out = String::new(); | ||
| while session.pos < max_pos { | ||
| let logits = transformer.forward(token, &mut session.kv_cache, session.pos as _); | ||
| let next = argmax(logits); | ||
|
|
||
| token = next; | ||
| session.pos += 1; | ||
| token = next; | ||
| session.pos += 1; | ||
|
|
||
| out.push_str(&self.tokenizer.decode(next).replace('▁', " ")); | ||
| } | ||
| out.push_str(&tokenizer.decode(next).replace('▁', " ")); | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,15 +1,5 @@ | ||
| use super::ServiceArgs; | ||
|
|
||
| pub(super) struct NvidiaService {} | ||
|
|
||
| impl From<ServiceArgs> for NvidiaService { | ||
| fn from(_: ServiceArgs) -> Self { | ||
| todo!() | ||
| } | ||
| } | ||
|
|
||
| impl NvidiaService { | ||
| pub fn run(self) { | ||
| todo!() | ||
| } | ||
| pub(super) fn run(args: ServiceArgs) { | ||
| todo!() | ||
| } |