import os import json import sys from argparse import ArgumentParser from typing import List import torch import torch.distributed as dist from transformers import AutoTokenizer from safetensors.torch import load_model from model import Transformer, ModelArgs current_dir = os.path.dirname(os.path.abspath(__file__)) encoding_dir = os.path.join(current_dir, "../encoding") sys.path.insert(0, os.path.abspath(encoding_dir)) from encoding_dsv4 import ( encode_case, load_cases, parse_message_from_completion_text, parse_tagged_text, ) from image_processor import prepare_vl_inputs @torch.inference_mode() def generate( model: Transformer, prompt_tokens: List[List[int]], max_new_tokens: int, eos_id: int, images=None, ) -> List[List[int]]: """Batch generation with right-padded prompts. The first forward pass processes [min_prompt_len:] tokens (prefill phase). Subsequent passes generate one token at a time (decode phase). For positions still within a prompt, the ground-truth token overrides the model's prediction. """ prompt_lens = [len(t) for t in prompt_tokens] assert max(prompt_lens) <= model.max_seq_len, f"Prompt length exceeds model maximum sequence length (max_seq_len={model.max_seq_len})" total_len = min(model.max_seq_len, max_new_tokens + max(prompt_lens)) tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long) for i, t in enumerate(prompt_tokens): tokens[i, :len(t)] = torch.tensor(t, dtype=torch.long) prev_pos = 0 finished = torch.tensor([False] * len(prompt_tokens)) prompt_mask = tokens != -1 for cur_pos in range(min(prompt_lens), total_len): next_token = model.forward(tokens[:, prev_pos:cur_pos], prev_pos, images)[0] next_token = torch.where(prompt_mask[:, cur_pos], tokens[:, cur_pos], next_token) tokens[:, cur_pos] = next_token finished |= torch.logical_and(~prompt_mask[:, cur_pos], next_token == eos_id) prev_pos = cur_pos if finished.all(): break completion_tokens = [] for i, toks in enumerate(tokens.tolist()): toks = toks[prompt_lens[i]:prompt_lens[i]+max_new_tokens] if eos_id in toks: toks = toks[:toks.index(eos_id)] toks.append(eos_id) completion_tokens.append(toks) return completion_tokens def generate_batched(model, prompt_tokens, images, max_new_tokens, eos_id, max_batch_size): """Prompts with images run alone (image spans need single-chunk prefill for in-image bidirectional attention); text-only prompts run in microbatches.""" if images is None: images = [None] * len(prompt_tokens) completion_tokens = [None] * len(prompt_tokens) for i, m in enumerate(images): if m is not None: completion_tokens[i] = generate(model, [prompt_tokens[i]], max_new_tokens, eos_id, [m])[0] text_idxs = [i for i, m in enumerate(images) if m is None] for j in range(0, len(text_idxs), max_batch_size): batch = text_idxs[j:j + max_batch_size] completions = generate(model, [prompt_tokens[i] for i in batch], max_new_tokens, eos_id) for i, toks in zip(batch, completions): completion_tokens[i] = toks return completion_tokens def prepare_case(case, thinking_mode, tokenizer, args): """Encode one message case and expand its image placeholders.""" if case.get("context"): raise ValueError( "Standalone inference does not support context without a prefilled KV cache") prompt, image_records = encode_case(case, thinking_mode) tokens, images = prepare_vl_inputs(prompt, image_records, tokenizer, args) return prompt, tokens, images def main( ckpt_path: str, config: str, input_file: str = "", interactive: bool = True, max_new_tokens: int = 100, temperature: float = 1.0, thinking_mode: str = "chat", ) -> None: world_size = int(os.getenv("WORLD_SIZE", "1")) rank = int(os.getenv("RANK", "0")) local_rank = int(os.getenv("LOCAL_RANK", "0")) if world_size > 1: dist.init_process_group("nccl") global print if rank != 0: print = lambda *_, **__: None torch.cuda.set_device(local_rank) torch.cuda.memory._set_allocator_settings("expandable_segments:True") torch.set_default_dtype(torch.bfloat16) torch.set_num_threads(8) torch.manual_seed(33377335) with open(config) as f: args = ModelArgs(**json.load(f)) args.temperature = temperature if interactive: args.max_batch_size = 1 args.max_seq_len = 64 * 1024 print(args) with torch.device("cuda"): model = Transformer(args) tokenizer = AutoTokenizer.from_pretrained(ckpt_path) print("load model") load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors")) torch.set_default_device("cuda") print("I'm DeepSeek 👋") if interactive: messages = [] while True: if world_size == 1: prompt = input(">>> ") elif rank == 0: prompt = input(">>> ") objects = [prompt] dist.broadcast_object_list(objects, 0) else: objects = [None] dist.broadcast_object_list(objects, 0) prompt = objects[0] if prompt == "/exit": break elif prompt == "/clear": messages.clear() continue messages.append({"role": "user", "content": parse_tagged_text(prompt)}) _, prompt_tokens, images = prepare_case( {"messages": messages}, thinking_mode, tokenizer, args) completion_tokens = generate(model, [prompt_tokens], max_new_tokens, tokenizer.eos_token_id, [images] if images else None) completion = tokenizer.decode(completion_tokens[0]) print(completion) messages.append(parse_message_from_completion_text(completion, thinking_mode=thinking_mode)) else: json_input = input_file.endswith(".json") if json_input: cases = load_cases(input_file) raw_prompts = None else: with open(input_file) as f: raw_prompts = f.read().rstrip("\n").split("\n\n") cases = [ {"messages": [{"role": "user", "content": parse_tagged_text(prompt)}]} for prompt in raw_prompts ] prepared = [prepare_case(case, thinking_mode, tokenizer, args) for case in cases] prompts = [prompt for prompt, _, _ in prepared] prompt_tokens = [tokens for _, tokens, _ in prepared] images = [image_inputs for _, _, image_inputs in prepared] if not any(images): images = None completion_tokens = generate_batched(model, prompt_tokens, images, max_new_tokens, tokenizer.eos_token_id, args.max_batch_size) for i, (case, prompt, toks) in enumerate(zip(cases, prompts, completion_tokens)): completion = tokenizer.decode(toks) print("Prompt:", prompt if json_input else raw_prompts[i]) print("Completion:", completion) if json_input: print("Parsed:", parse_message_from_completion_text( completion, thinking_mode=case.get("thinking_mode") or thinking_mode)) print() if world_size > 1: dist.destroy_process_group() if __name__ == "__main__": parser = ArgumentParser() parser.add_argument("--ckpt-path", type=str, required=True) parser.add_argument("--config", type=str, required=True) parser.add_argument("--input-file", type=str, default="", help="prompts separated by blank lines, or OpenAI-format cases (.json)") parser.add_argument("--interactive", action="store_true") parser.add_argument("--thinking-mode", type=str, default="chat", choices=["chat", "thinking"]) parser.add_argument("--max-new-tokens", type=int, default=None, help="default: 16384 for .json input, 200 otherwise") parser.add_argument("--temperature", type=float, default=None, help="default: 0.99 for .json input, 1.0 otherwise") args = parser.parse_args() assert args.input_file or args.interactive, "Either input-file or interactive mode must be specified" json_input = args.input_file.endswith(".json") if args.max_new_tokens is None: args.max_new_tokens = 16384 if json_input else 200 if args.temperature is None: args.temperature = 0.99 if json_input else 1.0 main(args.ckpt_path, args.config, args.input_file, args.interactive, args.max_new_tokens, args.temperature, args.thinking_mode)