model-conversion : use save_output_data for causual embeddings [no ci] (#26890)
This commit updates the python script that runs the original model to generate embeddings for the causal model, to use save_output_data which stores the token ids and the prompt in addition to logits. The motivation for this is that the embedding logits verification will fail as it expects these files (-prompt.txt and -tokens.bin) to exist. With the changes in this commit the causal-verify-embeddings target works again.
This commit is contained in:
@@ -2,12 +2,15 @@
|
|||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import os
|
import os
|
||||||
|
import sys
|
||||||
import importlib
|
import importlib
|
||||||
import torch
|
import torch
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
|
from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
|
||||||
from pathlib import Path
|
|
||||||
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
|
||||||
|
from utils.common import save_output_data
|
||||||
|
|
||||||
unreleased_model_name = os.getenv('UNRELEASED_MODEL_NAME')
|
unreleased_model_name = os.getenv('UNRELEASED_MODEL_NAME')
|
||||||
|
|
||||||
@@ -54,6 +57,7 @@ print(f"Model name: {model_name}")
|
|||||||
|
|
||||||
prompt = "Hello world today"
|
prompt = "Hello world today"
|
||||||
input_ids = tokenizer(prompt, return_tensors="pt").input_ids # ty: ignore[call-non-callable]
|
input_ids = tokenizer(prompt, return_tensors="pt").input_ids # ty: ignore[call-non-callable]
|
||||||
|
token_ids = input_ids[0].cpu().tolist()
|
||||||
print(f"Input tokens: {input_ids}")
|
print(f"Input tokens: {input_ids}")
|
||||||
print(f"Input text: {repr(prompt)}")
|
print(f"Input text: {repr(prompt)}")
|
||||||
print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}") # ty: ignore[unresolved-attribute]
|
print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}") # ty: ignore[unresolved-attribute]
|
||||||
@@ -74,21 +78,8 @@ with torch.no_grad():
|
|||||||
print(f"Hidden dimension: {token_embeddings.shape[-1]}")
|
print(f"Hidden dimension: {token_embeddings.shape[-1]}")
|
||||||
print(f"Number of tokens: {token_embeddings.shape[0]}")
|
print(f"Number of tokens: {token_embeddings.shape[0]}")
|
||||||
|
|
||||||
# Save raw token embeddings
|
|
||||||
data_dir = Path("data")
|
|
||||||
data_dir.mkdir(exist_ok=True)
|
|
||||||
bin_filename = data_dir / f"pytorch-{model_name}-embeddings.bin"
|
|
||||||
txt_filename = data_dir / f"pytorch-{model_name}-embeddings.txt"
|
|
||||||
|
|
||||||
# Save all token embeddings as binary
|
|
||||||
print(token_embeddings)
|
print(token_embeddings)
|
||||||
token_embeddings.astype(np.float32).tofile(bin_filename)
|
save_output_data(token_embeddings, token_ids, prompt, model_name, type_suffix="-embeddings")
|
||||||
|
|
||||||
# Save as text for inspection
|
|
||||||
with open(txt_filename, "w") as f:
|
|
||||||
for i, embedding in enumerate(token_embeddings):
|
|
||||||
for j, val in enumerate(embedding):
|
|
||||||
f.write(f"{i} {j} {val:.6f}\n")
|
|
||||||
|
|
||||||
# Print embeddings per token in the requested format
|
# Print embeddings per token in the requested format
|
||||||
print("\nToken embeddings:")
|
print("\nToken embeddings:")
|
||||||
@@ -110,5 +101,3 @@ with torch.no_grad():
|
|||||||
for i, token in enumerate(tokens):
|
for i, token in enumerate(tokens):
|
||||||
print(f" Token {i}: {repr(token)}")
|
print(f" Token {i}: {repr(token)}")
|
||||||
|
|
||||||
print(f"Saved bin logits to: {bin_filename}")
|
|
||||||
print(f"Saved txt logist to: {txt_filename}")
|
|
||||||
|
|||||||
Reference in New Issue
Block a user