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:
2026-05-21 22:04:00 +08:00
parent 6035ffdc0b
commit e1e75fc7f6
13 changed files with 971 additions and 0 deletions

View 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!();
}
}