char-fork/llama.cpp

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

Daniel Beveniusmodel-conversion : use save_output_data for causual embeddings [no ci] (#26890)1138b851f

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
4.1 KiB103 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, AutoConfig, AutoModelForCausalLM
11
12sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
13from utils.common import save_output_data
14
15unreleased_model_name = os.getenv('UNRELEASED_MODEL_NAME')
16
17parser = argparse.ArgumentParser(description='Process model with specified path')
18parser.add_argument('--model-path', '-m', help='Path to the model')
19args = parser.parse_args()
20
21model_path = os.environ.get('MODEL_PATH', args.model_path)
22if model_path is None:
23    parser.error("Model path must be specified either via --model-path argument or MODEL_PATH environment variable")
24
25config = AutoConfig.from_pretrained(model_path)
26
27print("Model type:       ", config.model_type)
28print("Vocab size:       ", config.vocab_size)
29print("Hidden size:      ", config.hidden_size)
30print("Number of layers: ", config.num_hidden_layers)
31print("BOS token id:     ", config.bos_token_id)
32print("EOS token id:     ", config.eos_token_id)
33
34print("Loading model and tokenizer using AutoTokenizer:", model_path)
35tokenizer = AutoTokenizer.from_pretrained(model_path)
36
37if unreleased_model_name:
38    model_name_lower = unreleased_model_name.lower()
39    unreleased_module_path = f"transformers.models.{model_name_lower}.modular_{model_name_lower}"
40    class_name = f"{unreleased_model_name}ForCausalLM"
41    print(f"Importing unreleased model module: {unreleased_module_path}")
42
43    try:
44        model_class = getattr(importlib.import_module(unreleased_module_path), class_name)
45        model = model_class.from_pretrained(model_path)
46    except (ImportError, AttributeError) as e:
47        print(f"Failed to import or load model: {e}")
48        print("Falling back to AutoModelForCausalLM")
49        model = AutoModelForCausalLM.from_pretrained(model_path)
50else:
51    model = AutoModelForCausalLM.from_pretrained(model_path)
52print(f"Model class: {type(model)}")
53#print(f"Model file: {type(model).__module__}")
54
55model_name = os.path.basename(model_path)
56print(f"Model name: {model_name}")
57
58prompt = "Hello world today"
59input_ids = tokenizer(prompt, return_tensors="pt").input_ids  # ty: ignore[call-non-callable]
60token_ids = input_ids[0].cpu().tolist()
61print(f"Input tokens: {input_ids}")
62print(f"Input text: {repr(prompt)}")
63print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}")  # ty: ignore[unresolved-attribute]
64
65with torch.no_grad():
66    outputs = model(input_ids, output_hidden_states=True)
67
68    # Extract hidden states from the last layer
69    # outputs.hidden_states is a tuple of (num_layers + 1) tensors
70    # Index -1 gets the last layer, shape: [batch_size, seq_len, hidden_size]
71    last_hidden_states = outputs.hidden_states[-1]
72
73    # Get embeddings for all tokens
74    token_embeddings = last_hidden_states[0].float().cpu().numpy()  # Remove batch dimension
75
76    print(f"Hidden states shape: {last_hidden_states.shape}")
77    print(f"Token embeddings shape: {token_embeddings.shape}")
78    print(f"Hidden dimension: {token_embeddings.shape[-1]}")
79    print(f"Number of tokens: {token_embeddings.shape[0]}")
80
81    print(token_embeddings)
82    save_output_data(token_embeddings, token_ids, prompt, model_name, type_suffix="-embeddings")
83
84    # Print embeddings per token in the requested format
85    print("\nToken embeddings:")
86    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])  # ty: ignore[unresolved-attribute]
87    for i, embedding in enumerate(token_embeddings):
88        # Format: show first few values, ..., then last few values
89        if len(embedding) > 10:
90            # Show first 3 and last 3 values with ... in between
91            first_vals = " ".join(f"{val:8.6f}" for val in embedding[:3])
92            last_vals = " ".join(f"{val:8.6f}" for val in embedding[-3:])
93            print(f"embedding {i}: {first_vals}  ... {last_vals}")
94        else:
95            # If embedding is short, show all values
96            vals = " ".join(f"{val:8.6f}" for val in embedding)
97            print(f"embedding {i}: {vals}")
98
99    # Also show token info for reference
100    print(f"\nToken reference:")
101    for i, token in enumerate(tokens):
102        print(f"  Token {i}: {repr(token)}")
103