| 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__)) |
| root_dir = os.path.abspath(os.path.join(current_dir, "..")) |
| encoding_dir = os.path.join(current_dir, '../encoding') |
| sys.path.insert(0, root_dir) |
| sys.path.insert(0, os.path.abspath(encoding_dir)) |
| from encoding_dsv4 import encode_messages, parse_message_from_completion_text |
| from cortex_ai.identity import project_system_prompt |
| from cortex_ai.vision.bridge import build_default_bridge |
|
|
|
|
| @torch.inference_mode() |
| def generate( |
| model: Transformer, |
| prompt_tokens: List[List[int]], |
| max_new_tokens: int, |
| eos_id: int, |
| ) -> List[List[int]]: |
| """Batch generation with left-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)[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 main( |
| ckpt_path: str, |
| config: str, |
| input_file: str = "", |
| interactive: bool = True, |
| max_new_tokens: int = 100, |
| temperature: float = 1.0, |
| vision_image: str = "", |
| ) -> 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"), strict=False) |
| torch.set_default_device("cuda") |
| print("I'm DeepSeek 👋") |
|
|
| if interactive: |
| if vision_image: |
| raise ValueError("--vision-image is supported in batch mode; use examples/vision_bridge.py for interactive observation") |
| messages = [{"role": "system", "content": project_system_prompt()}] |
| 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[:] = [{"role": "system", "content": project_system_prompt()}] |
| continue |
| messages.append({"role": "user", "content": prompt}) |
| prompt_tokens = tokenizer.encode(encode_messages(messages, thinking_mode="chat")) |
| completion_tokens = generate(model, [prompt_tokens], max_new_tokens, tokenizer.eos_token_id) |
| completion = tokenizer.decode(completion_tokens[0]) |
| print(completion) |
| messages.append(parse_message_from_completion_text(completion, thinking_mode="chat")) |
| else: |
| with open(input_file) as f: |
| prompts = f.read().split("\n\n") |
| vision_context = "" |
| if vision_image: |
| vision_context = build_default_bridge().build_context(vision_image) |
| prompt_tokens = [tokenizer.encode(encode_messages([ |
| {"role": "system", "content": project_system_prompt()}, |
| {"role": "user", "content": (vision_context + "\n\n" + prompt) if vision_context else prompt}, |
| ], thinking_mode="chat")) for prompt in prompts] |
| completion_tokens = generate(model, prompt_tokens, max_new_tokens, tokenizer.eos_token_id) |
| completions = tokenizer.batch_decode(completion_tokens) |
| for prompt, completion in zip(prompts, completions): |
| print("Prompt:", prompt) |
| print("Completion:", completion) |
| 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="") |
| parser.add_argument("--interactive", action="store_true") |
| parser.add_argument("--max-new-tokens", type=int, default=300) |
| parser.add_argument("--temperature", type=float, default=1.0) |
| parser.add_argument("--vision-image", type=str, default="", help="optional image path/URL for the external vision bridge in batch mode") |
| args = parser.parse_args() |
| assert args.input_file or args.interactive, "Either input-file or interactive mode must be specified" |
| main(args.ckpt_path, args.config, args.input_file, args.interactive, args.max_new_tokens, args.temperature, args.vision_image) |
|
|
