char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Daniel Beveniusmodel-conversion : make printing of config values optional (#19681)667b69427

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
6.3 KiB172 linesraw
1#!/usr/bin/env python3
2
3import argparse
4import os
5import sys
6import importlib
7import torch
8import numpy as np
9
10from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText, AutoConfig
11
12# Add parent directory to path for imports
13sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
14from utils.common import debug_hook, save_output_data
15
16def parse_arguments():
17    parser = argparse.ArgumentParser(description="Process model with specified path")
18    parser.add_argument("--model-path", "-m", help="Path to the model")
19    parser.add_argument("--prompt-file", "-f", help="Optional prompt file", required=False)
20    parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose debug output")
21    parser.add_argument("--device", "-d", help="Device to use (cpu, cuda, mps, auto)", default="auto")
22    return parser.parse_args()
23
24def load_model_and_tokenizer(model_path, device="auto"):
25    print("Loading model and tokenizer using AutoTokenizer:", model_path)
26    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
27    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
28    multimodal = False
29    full_config = config
30
31    # Determine device_map based on device argument
32    if device == "cpu":
33        device_map = {"": "cpu"}
34        print("Forcing CPU usage")
35    elif device == "auto":
36        device_map = "auto"
37    else:
38        device_map = {"": device}
39
40    print("Model type:       ", config.model_type)
41    if "vocab_size" not in config and "text_config" in config:
42        config = config.text_config
43        multimodal = True
44
45    def print_if_exists(label, obj, attr, default="N/A"):
46        val = getattr(obj, attr) if hasattr(obj, attr) else default
47        print(f"{label}", val)
48
49    print_if_exists("Vocab size:       ", config, "vocab_size")
50    print_if_exists("Hidden size:      ", config, "hidden_size")
51    print_if_exists("Number of layers: ", config, "num_hidden_layers")
52    print_if_exists("BOS token id:     ", config, "bos_token_id")
53    print_if_exists("EOS token id:     ", config, "eos_token_id")
54
55    unreleased_model_name = os.getenv("UNRELEASED_MODEL_NAME")
56    if unreleased_model_name:
57        model_name_lower = unreleased_model_name.lower()
58        unreleased_module_path = (
59            f"transformers.models.{model_name_lower}.modular_{model_name_lower}"
60        )
61        class_name = f"{unreleased_model_name}ForCausalLM"
62        print(f"Importing unreleased model module: {unreleased_module_path}")
63
64        try:
65            model_class = getattr(importlib.import_module(unreleased_module_path), class_name)
66            model = model_class.from_pretrained(
67                    model_path,
68                    device_map=device_map,
69                    offload_folder="offload",
70                    trust_remote_code=True,
71                    config=config
72            )
73        except (ImportError, AttributeError) as e:
74            print(f"Failed to import or load model: {e}")
75            exit(1)
76    else:
77        if multimodal:
78            model = AutoModelForImageTextToText.from_pretrained(
79                    model_path,
80                    device_map=device_map,
81                    offload_folder="offload",
82                    trust_remote_code=True,
83                    config=full_config
84            )
85        else:
86            model = AutoModelForCausalLM.from_pretrained(
87                    model_path,
88                    device_map=device_map,
89                    offload_folder="offload",
90                    trust_remote_code=True,
91                    config=config
92            )
93
94    print(f"Model class: {model.__class__.__name__}")
95
96    return model, tokenizer, config
97
98def enable_torch_debugging(model):
99        for name, module in model.named_modules():
100            if len(list(module.children())) == 0:  # only leaf modules
101                module.register_forward_hook(debug_hook(name))
102
103def get_prompt(args):
104    if args.prompt_file:
105        with open(args.prompt_file, encoding='utf-8') as f:
106            return f.read()
107    elif os.getenv("MODEL_TESTING_PROMPT"):
108        return os.getenv("MODEL_TESTING_PROMPT")
109    else:
110        return "Hello, my name is"
111
112def main():
113    args = parse_arguments()
114    model_path = os.environ.get("MODEL_PATH", args.model_path)
115    if model_path is None:
116        print("Error: Model path must be specified either via --model-path argument or MODEL_PATH environment variable")
117        sys.exit(1)
118
119
120    model, tokenizer, config = load_model_and_tokenizer(model_path, args.device)
121
122    if args.verbose:
123        enable_torch_debugging(model)
124
125    model_name = os.path.basename(model_path)
126
127    # Iterate over the model parameters (the tensors) and get the first one
128    # and use it to get the device the model is on.
129    device = next(model.parameters()).device
130    prompt = get_prompt(args)
131    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
132    token_ids = input_ids[0].cpu().tolist()
133
134    print(f"Input tokens: {input_ids}")
135    print(f"Input text: {repr(prompt)}")
136    print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}")
137
138    batch_size = 512
139
140    with torch.no_grad():
141        past = None
142        outputs = None
143        for i in range(0, input_ids.size(1), batch_size):
144            print(f"Processing chunk with tokens {i} to {i + batch_size}")
145            chunk = input_ids[:, i:i + batch_size]
146            outputs = model(chunk.to(model.device), past_key_values=past, use_cache=True)
147            past = outputs.past_key_values
148
149        logits = outputs.logits # type: ignore
150
151        # Extract logits for the last token (next token prediction)
152        last_logits = logits[0, -1, :].float().cpu().numpy()
153
154        print(f"Logits shape: {logits.shape}")
155        print(f"Last token logits shape: {last_logits.shape}")
156        print(f"Vocab size: {len(last_logits)}")
157
158        # Print some sample logits for quick verification
159        print(f"First 10 logits: {last_logits[:10]}")
160        print(f"Last 10 logits: {last_logits[-10:]}")
161
162        # Show top 5 predicted tokens
163        top_indices = np.argsort(last_logits)[-5:][::-1]
164        print("Top 5 predictions:")
165        for idx in top_indices:
166            token = tokenizer.decode([idx])
167            print(f"  Token {idx} ({repr(token)}): {last_logits[idx]:.6f}")
168
169        save_output_data(last_logits, token_ids, prompt, model_name)
170
171if __name__ == "__main__":
172    main()