style: format Rust workspace
This commit is contained in:
@@ -1,8 +1,8 @@
|
||||
use half::bf16;
|
||||
use std::path::PathBuf;
|
||||
use xserv_model::{loader, KVCache, ModelConfig, Qwen3};
|
||||
use xserv_model::{KVCache, ModelConfig, Qwen3, loader};
|
||||
use xserv_tensor::{DType, Device};
|
||||
use xserv_tokenizer::Tokenizer;
|
||||
use half::bf16;
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
@@ -20,8 +20,11 @@ fn main() {
|
||||
eprintln!("Token IDs: {token_ids:?}");
|
||||
|
||||
let mut cache = KVCache::new(
|
||||
config.num_layers(), config.num_kv_heads(), config.head_dim(),
|
||||
DType::BF16, Device::Cuda(0),
|
||||
config.num_layers(),
|
||||
config.num_kv_heads(),
|
||||
config.head_dim(),
|
||||
DType::BF16,
|
||||
Device::Cuda(0),
|
||||
);
|
||||
let logits = model.forward_with_cache(&token_ids, &mut cache);
|
||||
let logits_cpu = logits.to_device(Device::Cpu);
|
||||
@@ -31,7 +34,9 @@ fn main() {
|
||||
|
||||
// Print top-20 logits for the last position
|
||||
let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
|
||||
let mut indexed: Vec<(usize, f32)> = last_row.iter().enumerate()
|
||||
let mut indexed: Vec<(usize, f32)> = last_row
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| (i, v.to_f32()))
|
||||
.collect();
|
||||
indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
|
||||
Reference in New Issue
Block a user