| 52 | } |
| 53 | |
| 54 | static std::string category_for_tensor(const std::string & name) { |
| 55 | if (name == "causal_mask") return "causal_mask"; |
| 56 | if (name == "text_token_embed") return "text_token_embed"; |
| 57 | if (name == "text_after_pos_embed") return "text_after_pos_embed"; |
| 58 | if (name == "text_final_ln") return "text_final_ln"; |
| 59 | if (name == "text_features_2d") return "text_features_2d"; |
| 60 | |
| 61 | if (ends_with(name, "_after_ln1")) return "block_after_ln1"; |
| 62 | if (ends_with(name, "_qkv")) return "block_qkv"; |
| 63 | if (ends_with(name, "_attn_out")) return "block_attn_out"; |
| 64 | if (ends_with(name, "_after_attn_residual")) return "block_after_attn_residual"; |
| 65 | if (ends_with(name, "_after_ln2")) return "block_after_ln2"; |
| 66 | if (ends_with(name, "_mlp_fc1")) return "block_mlp_fc1"; |
| 67 | if (ends_with(name, "_mlp_gelu")) return "block_mlp_gelu"; |
| 68 | if (ends_with(name, "_mlp_out")) return "block_mlp_out"; |
| 69 | if (ends_with(name, "_out")) return "block_out"; |
| 70 | |
| 71 | return name; |
| 72 | } |
| 73 | |
| 74 | static const char * label_for_category(const std::string & category) { |
| 75 | if (category == "token_ids") return "Token IDs"; |