phase 6+7+8: model loading, BPE tokenizer, GPT-2 inference (Milestone ①)
Phase 6 — Model Loading (xserv-model): - safetensors parser with single/sharded file support - ModelConfig with dual naming (GPT-2 n_embd/n_head + modern HF naming) - Weight loading flow: safetensors → mmap → CPU Tensor → GPU Phase 7 — BPE Tokenizer (xserv-tokenizer): - Full BPE encode/decode from tokenizer.json - GPT-2 byte-to-unicode mapping (printable ASCII identity + shifted bytes) - Pre-tokenization regex, special token handling - Chat template support structure Phase 8 — GPT-2 Complete Inference: - GPT-2 model definition: wte, wpe, 12 transformer blocks, ln_f - Forward pass: embedding → (LayerNorm → MHA → residual → LayerNorm → MLP → residual) × 12 → LN → logits - QKV split with correct [batch, heads, seq, dim] layout (fixed reshape bug) - Greedy sampling from last-position logits - Interactive CLI: xserv-cli <model-dir> [--max-tokens N] Verified: GPT-2 124M generates coherent English text on RTX 5090. "The future of AI is uncertain. The future of AI is uncertain..." "Once upon a time, the world was a place of great beauty..." Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
78
crates/xserv-model/src/bin/xserv-cli.rs
Normal file
78
crates/xserv-model/src/bin/xserv-cli.rs
Normal file
@@ -0,0 +1,78 @@
|
||||
use std::io::{self, Write};
|
||||
use std::path::PathBuf;
|
||||
use xserv_model::{GPT2, ModelConfig};
|
||||
use xserv_model::loader;
|
||||
use xserv_model::gpt2::sample_greedy;
|
||||
use xserv_tokenizer::Tokenizer;
|
||||
use xserv_tensor::Device;
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
if args.len() < 2 {
|
||||
eprintln!("Usage: xserv-cli <model-dir> [--max-tokens N]");
|
||||
eprintln!(" model-dir: path to HF model directory (containing model.safetensors, config.json, tokenizer.json)");
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
let model_dir = PathBuf::from(&args[1]);
|
||||
let max_tokens: usize = args.iter()
|
||||
.position(|a| a == "--max-tokens")
|
||||
.and_then(|i| args.get(i + 1))
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(100);
|
||||
|
||||
xserv_cuda::device::set_device(0).unwrap();
|
||||
let info = xserv_cuda::device::device_info(0).unwrap();
|
||||
eprintln!("GPU: {} ({} MB free)", info.name, info.free_memory / 1024 / 1024);
|
||||
|
||||
// Load config
|
||||
let config = ModelConfig::from_file(&model_dir.join("config.json"));
|
||||
eprintln!("Model: {:?}, layers={}, hidden={}, heads={}, vocab={}",
|
||||
config.model_type, config.num_layers(), config.hidden(),
|
||||
config.num_heads(), config.vocab_size);
|
||||
|
||||
// Load weights
|
||||
eprintln!("Loading weights...");
|
||||
let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
|
||||
eprintln!("Loaded {} tensors", weights.len());
|
||||
|
||||
// GPT-2 uses weight names without "model." prefix
|
||||
let model = GPT2::from_weights(config, weights);
|
||||
|
||||
// Load tokenizer
|
||||
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
|
||||
eprintln!("Tokenizer loaded (vocab_size={})", tokenizer.vocab_size());
|
||||
eprintln!("Ready.\n");
|
||||
|
||||
// Interactive loop
|
||||
loop {
|
||||
print!("xserv> ");
|
||||
io::stdout().flush().unwrap();
|
||||
let mut input = String::new();
|
||||
if io::stdin().read_line(&mut input).unwrap() == 0 {
|
||||
break;
|
||||
}
|
||||
let input = input.trim();
|
||||
if input.is_empty() { continue; }
|
||||
if input == "quit" || input == "exit" { break; }
|
||||
|
||||
let mut token_ids = tokenizer.encode(input);
|
||||
print!("{input}");
|
||||
io::stdout().flush().unwrap();
|
||||
|
||||
for _ in 0..max_tokens {
|
||||
let logits = model.forward(&token_ids);
|
||||
let next = sample_greedy(&logits);
|
||||
token_ids.push(next);
|
||||
|
||||
let text = tokenizer.decode(&[next]);
|
||||
print!("{text}");
|
||||
io::stdout().flush().unwrap();
|
||||
|
||||
if tokenizer.eos_token_id() == Some(next) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
println!();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user