xserv/crates/xserv-model/src/bin/xserv-chat.rs

use std::io::{self, IsTerminal, Read, Write};
use std::path::PathBuf;

use xserv_model::{loader, sample, ModelConfig, PagedKVCache, Qwen3, SamplingParams, BLOCK_SIZE};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;

const SLOT: usize = 0;

struct CliOptions {
    model_dir: PathBuf,
    max_tokens: usize,
    max_seq_len: usize,
    sampling: SamplingParams,
    system_prompt: Option<String>,
    enable_thinking: bool,
    color: bool,
}

enum Finish {
    Stop { token_id: u32 },
    Length,
}

enum Line {
    Text(String),
    Eof,
}

/// RAII terminal raw-mode guard. Disables canonical mode + echo (keeps output
/// post-processing and signals), so we read keystrokes ourselves and edit the
/// line UTF-8-aware. Restores the original termios on drop.
struct RawMode {
    orig: libc::termios,
}

impl RawMode {
    fn enable() -> Option<Self> {
        unsafe {
            let mut orig: libc::termios = std::mem::zeroed();
            if libc::tcgetattr(libc::STDIN_FILENO, &mut orig) != 0 {
                return None;
            }
            let mut raw = orig;
            raw.c_lflag &= !(libc::ICANON | libc::ECHO);
            raw.c_cc[libc::VMIN as usize] = 1;
            raw.c_cc[libc::VTIME as usize] = 0;
            if libc::tcsetattr(libc::STDIN_FILENO, libc::TCSANOW, &raw) != 0 {
                return None;
            }
            Some(RawMode { orig })
        }
    }
}

impl Drop for RawMode {
    fn drop(&mut self) {
        unsafe {
            libc::tcsetattr(libc::STDIN_FILENO, libc::TCSANOW, &self.orig);
        }
    }
}

/// Read one line with UTF-8/CJK-aware editing. In a TTY this enters raw mode and
/// handles keystrokes so Backspace deletes a whole character (not a byte), and
/// multi-byte input (汉字/日本語/한글) renders correctly. Non-TTY (piped) input
/// falls back to a plain cooked read.
fn read_line_edited(prompt: &str) -> Line {
    let cooked = || -> Line {
        print!("{prompt}");
        io::stdout().flush().ok();
        let mut s = String::new();
        match io::stdin().read_line(&mut s) {
            Ok(0) | Err(_) => Line::Eof,
            Ok(_) => Line::Text(s),
        }
    };

    if !io::stdin().is_terminal() {
        return cooked();
    }
    let Some(_raw) = RawMode::enable() else {
        return cooked();
    };

    // Single-line editor: on every edit, rewrite the whole line so the terminal
    // renders correct (incl. double-width CJK) glyphs; \x1b[K clears leftovers.
    let redraw = |buf: &str| {
        print!("\r{prompt}{buf}\x1b[K");
        io::stdout().flush().ok();
    };

    let mut buf = String::new();
    redraw(&buf);
    let mut stdin = io::stdin().lock();
    let mut byte = [0u8; 1];

    loop {
        if stdin.read(&mut byte).unwrap_or(0) == 0 {
            // EOF on the stream.
            if buf.is_empty() {
                return Line::Eof;
            }
            break;
        }
        match byte[0] {
            b'\r' | b'\n' => {
                println!();
                break;
            }
            0x7f | 0x08 => {
                // Backspace: drop one whole char (String::pop is char-aware).
                buf.pop();
                redraw(&buf);
            }
            0x04 => {
                // Ctrl-D: EOF only when the line is empty.
                if buf.is_empty() {
                    return Line::Eof;
                }
            }
            0x1b => {
                // Escape sequence (arrows, etc.): consume and ignore the 2 bytes
                // of a typical CSI sequence so they don't land in the buffer.
                let mut seq = [0u8; 2];
                let _ = stdin.read(&mut seq);
            }
            b if b < 0x20 => { /* other control bytes: ignore */ }
            b if b < 0x80 => {
                buf.push(b as char);
                redraw(&buf);
            }
            b => {
                // UTF-8 multi-byte: read the continuation bytes for this char.
                let extra = if b >= 0xF0 { 3 } else if b >= 0xE0 { 2 } else { 1 };
                let mut bytes = vec![b];
                let mut cont = [0u8; 1];
                let mut ok = true;
                for _ in 0..extra {
                    if stdin.read(&mut cont).unwrap_or(0) == 0 {
                        ok = false;
                        break;
                    }
                    bytes.push(cont[0]);
                }
                if ok {
                    if let Ok(s) = std::str::from_utf8(&bytes) {
                        buf.push_str(s);
                        redraw(&buf);
                    }
                }
            }
        }
    }
    Line::Text(buf)
}

fn main() {
    let opts = parse_args();

    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
    );

    let config = ModelConfig::from_file(&opts.model_dir.join("config.json"));
    let model_type = config.model_type.as_deref().unwrap_or("unknown");
    if !model_type.contains("qwen") {
        eprintln!("xserv-chat currently supports Qwen-style ChatML models only; got model_type={model_type}");
        std::process::exit(2);
    }

    let max_seq_len = opts.max_seq_len.min(config.max_seq_len()).max(1);
    eprintln!(
        "Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}, max_seq_len={}",
        config.num_layers(),
        config.hidden(),
        config.num_heads(),
        config.num_kv_heads(),
        config.vocab_size,
        max_seq_len
    );

    eprintln!("Loading weights...");
    let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0));
    eprintln!("Loaded {} tensors", weights.len());
    let model = Qwen3::from_weights(config.clone(), weights);

    let tokenizer = Tokenizer::from_file(&opts.model_dir.join("tokenizer.json"));
    let mut cache = new_paged_cache(&config, max_seq_len);
    cache.register_sequence(SLOT).expect("register chat slot");
    let use_color = opts.color && io::stdout().is_terminal();

    eprintln!("Ready (paged KV cache, persistent chat slot).");
    eprintln!("Commands: /exit, /quit, /clear\n");

    loop {
        let line = match read_line_edited("user> ") {
            Line::Eof => break,
            Line::Text(s) => s,
        };
        let input = line.trim();
        if input.is_empty() {
            continue;
        }

        match input {
            "/exit" | "/quit" | "exit" | "quit" => break,
            "/clear" => {
                cache.free_sequence(SLOT);
                cache.register_sequence(SLOT).expect("register chat slot");
                eprintln!("history and KV cache cleared");
                continue;
            }
            "/help" => {
                print_help();
                continue;
            }
            _ => {}
        }

        let include_system = cache.seq_len(SLOT) == 0;
        let prompt = build_turn_prompt(
            opts.system_prompt.as_deref(),
            include_system,
            input,
            opts.enable_thinking,
        );
        let prompt_tokens = tokenizer.encode(&prompt);
        if prompt_tokens.is_empty() {
            continue;
        }

        let used = cache.seq_len(SLOT);
        let remaining = max_seq_len.saturating_sub(used);
        if prompt_tokens.len() >= remaining {
            eprintln!(
                "context full: {used}/{max_seq_len} tokens used, new turn needs {} tokens; use /clear",
                prompt_tokens.len()
            );
            continue;
        }
        let max_new_tokens = opts.max_tokens.min(remaining - prompt_tokens.len());

        print!("assistant> ");
        io::stdout().flush().unwrap();
        let finish = generate_with_paged_cache(
            &model,
            &mut cache,
            &tokenizer,
            &prompt_tokens,
            &opts.sampling,
            max_new_tokens,
            use_color,
        );
        match finish {
            Finish::Stop { token_id } => {
                append_after_stop(&model, &mut cache, &tokenizer, max_seq_len, token_id);
            }
            Finish::Length => {
                append_text_to_cache(&model, &mut cache, &tokenizer, max_seq_len, "<|im_end|>\n");
            }
        }
        println!();
    }
}

fn parse_args() -> CliOptions {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.is_empty() || args.iter().any(|a| a == "--help" || a == "-h") {
        print_usage_and_exit(0);
    }

    let mut model_dir = None;
    let mut max_tokens = 256usize;
    let mut max_seq_len = 2048usize;
    let mut temperature = 0.0f32;
    let mut top_k = 0usize;
    let mut top_p = 1.0f32;
    let mut system_prompt = None;
    let mut enable_thinking = false;
    let mut color = true;

    let mut i = 0;
    while i < args.len() {
        match args[i].as_str() {
            "-m" | "--model" => {
                i += 1;
                model_dir = args.get(i).map(PathBuf::from);
            }
            "--max-tokens" => {
                i += 1;
                max_tokens = parse_value(&args, i, "--max-tokens");
            }
            "--max-seq-len" => {
                i += 1;
                max_seq_len = parse_value(&args, i, "--max-seq-len");
            }
            "--temperature" => {
                i += 1;
                temperature = parse_value(&args, i, "--temperature");
            }
            "--top-k" => {
                i += 1;
                top_k = parse_value(&args, i, "--top-k");
            }
            "--top-p" => {
                i += 1;
                top_p = parse_value(&args, i, "--top-p");
            }
            "--system" => {
                i += 1;
                system_prompt = args.get(i).cloned();
                if system_prompt.is_none() {
                    eprintln!("missing value for --system");
                    std::process::exit(2);
                }
            }
            "--think" => {
                enable_thinking = true;
            }
            "--no-color" => {
                color = false;
            }
            arg if arg.starts_with('-') => {
                eprintln!("unknown option: {arg}");
                print_usage_and_exit(2);
            }
            arg => {
                if model_dir.is_some() {
                    eprintln!("unexpected extra argument: {arg}");
                    print_usage_and_exit(2);
                }
                model_dir = Some(PathBuf::from(arg));
            }
        }
        i += 1;
    }

    CliOptions {
        model_dir: model_dir.unwrap_or_else(|| {
            eprintln!("missing model directory");
            print_usage_and_exit(2);
        }),
        max_tokens: max_tokens.max(1),
        max_seq_len: max_seq_len.max(1),
        sampling: SamplingParams {
            temperature,
            top_k,
            top_p,
        },
        system_prompt,
        enable_thinking,
        color,
    }
}

fn parse_value<T: std::str::FromStr>(args: &[String], i: usize, name: &str) -> T {
    args.get(i).and_then(|s| s.parse().ok()).unwrap_or_else(|| {
        eprintln!("invalid or missing value for {name}");
        std::process::exit(2);
    })
}

fn print_usage_and_exit(code: i32) -> ! {
    eprintln!(
        "Usage: xserv-chat <model-dir> [options]\n\
         \n\
         Options:\n\
         \t-m, --model DIR       Model directory\n\
         \t--max-tokens N        Max generated tokens per turn (default: 256)\n\
         \t--max-seq-len N       Persistent KV context length (default: 2048)\n\
         \t--temperature F       Sampling temperature, 0 = greedy (default: 0)\n\
         \t--top-k N             Top-k sampling, 0 = disabled (default: 0)\n\
         \t--top-p F             Top-p sampling (default: 1.0)\n\
         \t--system TEXT         System prompt for the first turn after start or /clear\n\
         \t--think               Let Qwen3 emit thinking; rendered gray on terminals\n\
         \t--no-color            Disable ANSI color for thinking output\n\
         \t-h, --help            Show this help"
    );
    std::process::exit(code);
}

fn print_help() {
    eprintln!("Commands:");
    eprintln!("  /clear  clear chat history and free/recreate the paged KV slot");
    eprintln!("  /exit   quit");
    eprintln!("  /quit   quit");
}

fn new_paged_cache(config: &ModelConfig, max_seq_len: usize) -> PagedKVCache {
    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
    let total_blocks = (max_blocks_per_seq + 1).max(2);
    // Single-slot interactive CLI: no swap pool (cpu_total_blocks = 0).
    PagedKVCache::new(config, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, 0)
}

fn build_turn_prompt(
    system: Option<&str>,
    include_system: bool,
    user_input: &str,
    enable_thinking: bool,
) -> String {
    let mut prompt = String::new();
    if include_system {
        if let Some(system) = system {
            if !system.trim().is_empty() {
                prompt.push_str("<|im_start|>system\n");
                prompt.push_str(system.trim());
                prompt.push_str("<|im_end|>\n");
            }
        }
    }
    prompt.push_str("<|im_start|>user\n");
    prompt.push_str(user_input);
    prompt.push_str("<|im_end|>\n");
    prompt.push_str("<|im_start|>assistant\n");
    if !enable_thinking {
        prompt.push_str("<think>\n\n</think>\n\n");
    }
    prompt
}

fn generate_with_paged_cache(
    model: &Qwen3,
    cache: &mut PagedKVCache,
    tokenizer: &Tokenizer,
    prompt_tokens: &[u32],
    sampling: &SamplingParams,
    max_tokens: usize,
    use_color: bool,
) -> Finish {
    let logits = model.forward_prefill_paged(prompt_tokens, SLOT, cache);
    let mut next = sample(&logits, sampling);
    let mut decode_buffer = Vec::new();
    let mut in_thinking = false;

    for _ in 0..max_tokens {
        let position = cache.seq_len(SLOT);
        let logits = model.forward_decode_paged(&[next], &[position], &[SLOT], cache);
        if is_stop_token(tokenizer, next) {
            print_stream_text(
                &tokenizer.flush_decode_stream(&mut decode_buffer),
                in_thinking,
                use_color,
            );
            io::stdout().flush().unwrap();
            return Finish::Stop { token_id: next };
        }

        print_generated_token(
            tokenizer,
            next,
            &mut decode_buffer,
            &mut in_thinking,
            use_color,
        );
        io::stdout().flush().unwrap();
        next = sample(&logits, sampling);
    }

    print_stream_text(
        &tokenizer.flush_decode_stream(&mut decode_buffer),
        in_thinking,
        use_color,
    );
    io::stdout().flush().unwrap();
    Finish::Length
}

fn append_after_stop(
    model: &Qwen3,
    cache: &mut PagedKVCache,
    tokenizer: &Tokenizer,
    max_seq_len: usize,
    stop_token_id: u32,
) {
    if tokenizer.special_token_id("<|im_end|>") == Some(stop_token_id) {
        append_text_to_cache(model, cache, tokenizer, max_seq_len, "\n");
    }
}

fn append_text_to_cache(
    model: &Qwen3,
    cache: &mut PagedKVCache,
    tokenizer: &Tokenizer,
    max_seq_len: usize,
    text: &str,
) {
    let tokens = tokenizer.encode(text);
    if tokens.is_empty() || cache.seq_len(SLOT) + tokens.len() > max_seq_len {
        return;
    }
    let _ = model.forward_prefill_paged(&tokens, SLOT, cache);
}

fn print_generated_token(
    tokenizer: &Tokenizer,
    token_id: u32,
    decode_buffer: &mut Vec<u8>,
    in_thinking: &mut bool,
    use_color: bool,
) {
    if tokenizer.special_token_id("<think>") == Some(token_id) {
        print_stream_text(
            &tokenizer.flush_decode_stream(decode_buffer),
            *in_thinking,
            use_color,
        );
        *in_thinking = true;
        print_stream_text("<think>", true, use_color);
        return;
    }

    if tokenizer.special_token_id("</think>") == Some(token_id) {
        print_stream_text(
            &tokenizer.flush_decode_stream(decode_buffer),
            *in_thinking,
            use_color,
        );
        print_stream_text("</think>", true, use_color);
        *in_thinking = false;
        return;
    }

    let text = tokenizer.decode_token_stream(token_id, decode_buffer);
    print_stream_text(&text, *in_thinking, use_color);
}

fn print_stream_text(text: &str, in_thinking: bool, use_color: bool) {
    if text.is_empty() {
        return;
    }
    if in_thinking && use_color {
        print!("\x1b[90m{text}\x1b[0m");
    } else {
        print!("{text}");
    }
}

fn is_stop_token(tokenizer: &Tokenizer, token_id: u32) -> bool {
    tokenizer.eos_token_id() == Some(token_id)
        || tokenizer.special_token_id("<|im_end|>") == Some(token_id)
        || tokenizer.special_token_id("<|endoftext|>") == Some(token_id)
        || tokenizer.special_token_id("<|end_of_text|>") == Some(token_id)
}