()
| 94 | |
| 95 | |
| 96 | def main() -> None: |
| 97 | parser = argparse.ArgumentParser() |
| 98 | parser.add_argument("--checkpoint", required=True) |
| 99 | parser.add_argument("--tokenizer", default="raw_weights/tokenizer.json") |
| 100 | parser.add_argument("--prompts", default="tests/phase4_prompts.tsv") |
| 101 | parser.add_argument("--outdir", default="tests/ref_phase4") |
| 102 | args = parser.parse_args() |
| 103 | |
| 104 | os.makedirs(args.outdir, exist_ok=True) |
| 105 | device = "cpu" |
| 106 | |
| 107 | print("Loading checkpoint...") |
| 108 | ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) |
| 109 | if "model" in ckpt: |
| 110 | ckpt = ckpt["model"] |
| 111 | |
| 112 | text_prefix = "detector.backbone.language_backbone.encoder." |
| 113 | text_sd = {k[len(text_prefix):]: v for k, v in ckpt.items() if k.startswith(text_prefix)} |
| 114 | resizer_prefix = "detector.backbone.language_backbone.resizer." |
| 115 | resizer_sd = { |
| 116 | k[len(resizer_prefix):]: v for k, v in ckpt.items() if k.startswith(resizer_prefix) |
| 117 | } |
| 118 | |
| 119 | tokenizer = Tokenizer.from_file(args.tokenizer) |
| 120 | prompts = load_prompt_cases(args.prompts) |
| 121 | |
| 122 | text_width = 1024 |
| 123 | text_heads = 16 |
| 124 | text_layers = 24 |
| 125 | context_length = 32 |
| 126 | |
| 127 | causal_mask = torch.full((context_length, context_length), float("-inf"), device=device) |
| 128 | causal_mask = torch.triu(causal_mask, diagonal=1) |
| 129 | |
| 130 | with torch.no_grad(): |
| 131 | for prompt_id, prompt_text in prompts: |
| 132 | prompt_dir = os.path.join(args.outdir, prompt_id) |
| 133 | os.makedirs(prompt_dir, exist_ok=True) |
| 134 | |
| 135 | token_ids = tokenize_prompt(tokenizer, prompt_text, context_length) |
| 136 | with open(os.path.join(prompt_dir, "prompt.txt"), "w", encoding="utf-8") as f: |
| 137 | f.write(prompt_text) |
| 138 | |
| 139 | save_i32_tensor(os.path.join(prompt_dir, "token_ids"), token_ids) |
| 140 | save_tensor(os.path.join(prompt_dir, "causal_mask"), causal_mask) |
| 141 | |
| 142 | tokens = torch.tensor([token_ids], dtype=torch.long, device=device) |
| 143 | |
| 144 | x_bf = F.embedding(tokens, text_sd["token_embedding.weight"]) |
| 145 | save_ggml_2d_batch_first(os.path.join(prompt_dir, "text_token_embed"), x_bf) |
| 146 | |
| 147 | x_bf = x_bf + text_sd["positional_embedding"][:context_length] |
| 148 | save_ggml_2d_batch_first(os.path.join(prompt_dir, "text_after_pos_embed"), x_bf) |
| 149 | |
| 150 | x = x_bf.transpose(0, 1).contiguous() # [T, 1, D] |
| 151 | |
| 152 | mha = nn.MultiheadAttention(text_width, text_heads, batch_first=False).to(device) |
| 153 | mha.eval() |
no test coverage detected