Felipe97/llama-cpp-compiled

1#!/usr/bin/env python32 3import argparse4import os5import sys6import importlib7 8from transformers import AutoTokenizer, AutoConfig, AutoModel9import torch10 11# Add parent directory to path for imports12sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))13from utils.common import save_output_data14 15 16def parse_arguments():17    parser = argparse.ArgumentParser(description='Run original embedding model')18    parser.add_argument(19        '--model-path',20        '-m',21        help='Path to the model'22    )23    parser.add_argument(24        '--prompts-file',25        '-p',26        help='Path to file containing prompts (one per line)'27    )28    parser.add_argument(29        '--use-sentence-transformers',30        action='store_true',31        help=('Use SentenceTransformer to apply all numbered layers '32              '(01_Pooling, 02_Dense, 03_Dense, 04_Normalize)')33    )34    parser.add_argument(35        '--device',36        '-d',37        help='Device to use (cpu, cuda, mps, auto)',38        default='auto'39    )40    return parser.parse_args()41 42 43def load_model_and_tokenizer(model_path, use_sentence_transformers=False, device="auto"):44    if device == "cpu":45        device_map = {"": "cpu"}46        print("Forcing CPU usage")47    elif device == "auto":48        # On Mac, "auto" device_map can cause issues with accelerate49        # So we detect the best device manually50        if torch.cuda.is_available():51            device_map = {"": "cuda"}52            print("Using CUDA")53        elif torch.backends.mps.is_available():54            device_map = {"": "mps"}55            print("Using MPS (Apple Metal)")56        else:57            device_map = {"": "cpu"}58            print("Using CPU")59    else:60        device_map = {"": device}61 62    if use_sentence_transformers:63        from sentence_transformers import SentenceTransformer64        print("Using SentenceTransformer to apply all numbered layers")65        model = SentenceTransformer(model_path)66        tokenizer = model.tokenizer67        config = model[0].auto_model.config  # ty: ignore[unresolved-attribute]68    else:69        tokenizer = AutoTokenizer.from_pretrained(model_path)70        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)71 72        # This can be used to override the sliding window size for manual testing. This73        # can be useful to verify the sliding window attention mask in the original model74        # and compare it with the converted .gguf model.75        if hasattr(config, 'sliding_window'):76            original_sliding_window = config.sliding_window77            print(f"Modified sliding window: {original_sliding_window} -> {config.sliding_window}")78 79        unreleased_model_name = os.getenv('UNRELEASED_MODEL_NAME')80        print(f"Using unreleased model: {unreleased_model_name}")81        if unreleased_model_name:82            model_name_lower = unreleased_model_name.lower()83            unreleased_module_path = f"transformers.models.{model_name_lower}.modular_{model_name_lower}"84            class_name = f"{unreleased_model_name}Model"85            print(f"Importing unreleased model module: {unreleased_module_path}")86 87            try:88                model_class = getattr(importlib.import_module(unreleased_module_path), class_name)89                model = model_class.from_pretrained(90                    model_path,91                    device_map=device_map,92                    offload_folder="offload",93                    trust_remote_code=True,94                    config=config95                )96            except (ImportError, AttributeError) as e:97                print(f"Failed to import or load model: {e}")98                sys.exit(1)99        else:100            model = AutoModel.from_pretrained(101                model_path,102                device_map=device_map,103                offload_folder="offload",104                trust_remote_code=True,105                config=config106            )107        print(f"Model class: {type(model)}")108        print(f"Model file: {type(model).__module__}")109 110        # Verify the model is using the correct sliding window111        if hasattr(model.config, 'sliding_window'):112            print(f"Model's sliding_window: {model.config.sliding_window}")113        else:114            print("Model config does not have sliding_window attribute")115 116    return model, tokenizer, config117 118 119def get_prompt(args):120    if args.prompts_file:121        try:122            with open(args.prompts_file, 'r', encoding='utf-8') as f:123                return f.read().strip()124        except FileNotFoundError:125            print(f"Error: Prompts file '{args.prompts_file}' not found")126            sys.exit(1)127        except Exception as e:128            print(f"Error reading prompts file: {e}")129            sys.exit(1)130    else:131        return "Hello world today"132 133 134def main():135    args = parse_arguments()136 137    model_path = os.environ.get('EMBEDDING_MODEL_PATH', args.model_path)138    if model_path is None:139        print("Error: Model path must be specified either via --model-path argument "140              "or EMBEDDING_MODEL_PATH environment variable")141        sys.exit(1)142 143    # Determine if we should use SentenceTransformer144    use_st = (145        args.use_sentence_transformers or os.environ.get('USE_SENTENCE_TRANSFORMERS', '').lower() in ('1', 'true', 'yes')146    )147 148    model, tokenizer, config = load_model_and_tokenizer(model_path, use_st, args.device)149 150    # Get the device the model is on151    if not use_st:152        device = next(model.parameters()).device153    else:154        # For SentenceTransformer, get device from the underlying model155        device = next(model[0].auto_model.parameters()).device156 157    model_name = os.path.basename(model_path)158 159    prompt_text = get_prompt(args)160    texts = [prompt_text]161 162    with torch.no_grad():163        if use_st:164            embeddings = model.encode(texts, convert_to_numpy=True)165            all_embeddings = embeddings  # Shape: [batch_size, hidden_size]166 167            encoded = tokenizer(168                texts,169                padding=True,170                truncation=True,171                return_tensors="pt"172            )173            tokens = encoded['input_ids'][0]174            token_ids = tokens.cpu().tolist()175            token_strings = tokenizer.convert_ids_to_tokens(tokens)176            for i, (token_id, token_str) in enumerate(zip(tokens, token_strings)):177                print(f"{token_id:6d} -> '{token_str}'")178 179            print(f"Embeddings shape (after all SentenceTransformer layers): {all_embeddings.shape}")180            print(f"Embedding dimension: {all_embeddings.shape[1] if len(all_embeddings.shape) > 1 else all_embeddings.shape[0]}")181        else:182            # Standard approach: use base model output only183            encoded = tokenizer(184                texts,185                padding=True,186                truncation=True,187                return_tensors="pt"188            )189 190            tokens = encoded['input_ids'][0]191            token_ids = tokens.cpu().tolist()192            token_strings = tokenizer.convert_ids_to_tokens(tokens)193            for i, (token_id, token_str) in enumerate(zip(tokens, token_strings)):194                print(f"{token_id:6d} -> '{token_str}'")195 196            # Move inputs to the same device as the model197            encoded = {k: v.to(device) for k, v in encoded.items()}198            outputs = model(**encoded)199            hidden_states = outputs.last_hidden_state  # Shape: [batch_size, seq_len, hidden_size]200 201            all_embeddings = hidden_states[0].float().cpu().numpy()  # Shape: [seq_len, hidden_size]202 203            print(f"Hidden states shape: {hidden_states.shape}")204            print(f"All embeddings shape: {all_embeddings.shape}")205            print(f"Embedding dimension: {all_embeddings.shape[1]}")206 207        if len(all_embeddings.shape) == 1:208            n_embd = all_embeddings.shape[0]209            n_embd_count = 1210            all_embeddings = all_embeddings.reshape(1, -1)211        else:212            n_embd = all_embeddings.shape[1]213            n_embd_count = all_embeddings.shape[0]214 215        print()216 217        for j in range(n_embd_count):218            embedding = all_embeddings[j]219            print(f"embedding {j}: ", end="")220 221            # Print first 3 values222            for i in range(min(3, n_embd)):223                print(f"{embedding[i]:9.6f} ", end="")224 225            print(" ... ", end="")226 227            # Print last 3 values228            for i in range(n_embd - 3, n_embd):229                print(f"{embedding[i]:9.6f} ", end="")230 231            print()  # New line232 233        print()234 235        flattened_embeddings = all_embeddings.flatten()236        print(f"Total values: {len(flattened_embeddings)} ({n_embd_count} embeddings × {n_embd} dimensions)")237        print("")238 239        save_output_data(flattened_embeddings, token_ids, prompt_text, model_name, type_suffix="-embeddings")240 241 242if __name__ == "__main__":243    main()244