Felipe97/llama-cpp-compiled

1#!/usr/bin/env python32 3import argparse4import os5import sys6import importlib7import torch8import numpy as np9 10from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText, AutoConfig11 12# Add parent directory to path for imports13sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))14from utils.common import debug_hook, save_output_data15 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 = False29    full_config = config30 31    # Determine device_map based on device argument32    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_config43        multimodal = True44 45    def print_if_exists(label, obj, attr, default="N/A"):46        val = getattr(obj, attr) if hasattr(obj, attr) else default47        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=config72            )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_config84            )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=config92            )93 94    print(f"Model class: {model.__class__.__name__}")95 96    return model, tokenizer, config97 98def enable_torch_debugging(model):99        for name, module in model.named_modules():100            if len(list(module.children())) == 0:  # only leaf modules101                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 one128    # and use it to get the device the model is on.129    device = next(model.parameters()).device130    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 = 512139 140    with torch.no_grad():141        past = None142        outputs = None143        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_values148 149        logits = outputs.logits # type: ignore150 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 verification159        print(f"First 10 logits: {last_logits[:10]}")160        print(f"Last 10 logits: {last_logits[-10:]}")161 162        # Show top 5 predicted tokens163        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()173