145 lines
6.2 KiB
Python
145 lines
6.2 KiB
Python
import os
|
||
import argparse
|
||
import torch
|
||
from transformers import AutoTokenizer
|
||
from model.model import MiniMindLM, ExtractDB
|
||
from model.LMConfig import LMConfig
|
||
|
||
def decode_dataset(model_path, output_path, device="cuda"):
|
||
"""
|
||
Decode the weight_down_embed buffer in the model to readable text
|
||
|
||
Args:
|
||
model_path: Path to the model checkpoint
|
||
output_path: Path to save the decoded text
|
||
device: Device to load the model on
|
||
"""
|
||
print(f"Loading tokenizer from ./model/minimind_tokenizer")
|
||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||
|
||
print(f"Setting up model configuration")
|
||
# Create model configuration matching the training parameters
|
||
lm_config = LMConfig(
|
||
dim=1024,
|
||
n_layers=32,
|
||
max_seq_len=1024,
|
||
use_flash_attn=True,
|
||
knowledge_num=16384, # From the script parameters
|
||
knowledge_length=64 # From the script parameters
|
||
)
|
||
|
||
print(f"Initializing model")
|
||
model = MiniMindLM(lm_config).to(device)
|
||
|
||
print(f"Loading model weights from {model_path}")
|
||
state_dict = torch.load(model_path, map_location=device)
|
||
|
||
# Get model parameters
|
||
model_state = dict(model.named_parameters())
|
||
model_state.update(dict(model.named_buffers()))
|
||
|
||
# Find parameters with matching names but different shapes
|
||
shape_mismatch = {}
|
||
for name, param in model_state.items():
|
||
if name in state_dict and param.shape != state_dict[name].shape:
|
||
shape_mismatch[name] = (param.shape, state_dict[name].shape)
|
||
|
||
# Find parameters in model but not in state_dict and vice versa
|
||
model_only = set(model_state.keys()) - set(state_dict.keys())
|
||
state_dict_only = set(state_dict.keys()) - set(model_state.keys())
|
||
|
||
# Create filtered state_dict with only compatible parameters
|
||
filtered_state_dict = {}
|
||
for name, param in state_dict.items():
|
||
if name in model_state and param.shape == model_state[name].shape:
|
||
filtered_state_dict[name] = param
|
||
|
||
# Print parameter differences
|
||
if shape_mismatch:
|
||
print(f"Parameters with shape mismatches: {len(shape_mismatch)}")
|
||
for name, (model_shape, state_shape) in shape_mismatch.items():
|
||
print(f" {name}: model={model_shape}, checkpoint={state_shape}")
|
||
|
||
if model_only:
|
||
print(f"Parameters in model but not in checkpoint: {len(model_only)}")
|
||
for name in sorted(model_only):
|
||
print(f" {name}: {model_state[name].shape}")
|
||
|
||
# 特殊处理pos_cis_real参数
|
||
if name == "pos_cis_real":
|
||
print(f"Detected pos_cis_real parameter. This is a position encoding that will be initialized automatically.")
|
||
|
||
if state_dict_only:
|
||
print(f"Parameters in checkpoint but not in model: {len(state_dict_only)}")
|
||
for name in sorted(state_dict_only):
|
||
print(f" {name}: {state_dict[name].shape}")
|
||
|
||
# 如果checkpoint中有output.weight但模型中没有，尝试加载到tok_embeddings
|
||
if name == "output.weight" and "tok_embeddings.weight" in model_state:
|
||
print(f"Found output.weight in checkpoint but not in model. Will try to map it to tok_embeddings.weight")
|
||
if model_state["tok_embeddings.weight"].shape == state_dict["output.weight"].shape:
|
||
filtered_state_dict["tok_embeddings.weight"] = state_dict["output.weight"]
|
||
|
||
# Load only the compatible parameters
|
||
print(f"Loading {len(filtered_state_dict)}/{len(state_dict)} parameters")
|
||
model.load_state_dict(filtered_state_dict, strict=False)
|
||
|
||
# 检查extract_db和weight_down_embed是否存在
|
||
if not hasattr(model, "extract_db"):
|
||
print("ERROR: Model does not have extract_db attribute. This is required for decoding.")
|
||
return
|
||
|
||
print("Accessing weight_down_embed buffer")
|
||
# Get the weight_down_embed buffer from the model
|
||
try:
|
||
weight_down_embed = model.extract_db.weight_down_embed
|
||
print(f"Successfully accessed weight_down_embed buffer")
|
||
except Exception as e:
|
||
print(f"ERROR: Failed to access weight_down_embed buffer: {e}")
|
||
print(f"Model structure: {model.__class__.__name__}")
|
||
print(f"ExtractDB attributes: {dir(model.extract_db)}")
|
||
return
|
||
|
||
print(f"Shape of weight_down_embed: {weight_down_embed.shape}")
|
||
print(f"Data type of weight_down_embed: {weight_down_embed.dtype}")
|
||
|
||
# Create output directory if it doesn't exist
|
||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||
|
||
print(f"Decoding knowledge and writing to {output_path}")
|
||
knowledge_num, knowledge_length = weight_down_embed.shape
|
||
|
||
with open(output_path, 'w', encoding='utf-8') as f:
|
||
for i in range(knowledge_num):
|
||
try:
|
||
# Get token IDs for this knowledge entry
|
||
token_ids = weight_down_embed[i].cpu().tolist()
|
||
|
||
# Decode tokens to text
|
||
text = tokenizer.decode(token_ids, skip_special_tokens=True)
|
||
|
||
# Write to file
|
||
f.write(f"Knowledge_{i}: {text}\n")
|
||
|
||
# Print progress periodically
|
||
if (i + 1) % 100 == 0:
|
||
print(f"Decoded {i + 1}/{knowledge_num} knowledge entries")
|
||
except Exception as e:
|
||
print(f"Error decoding knowledge entry {i}: {e}")
|
||
f.write(f"Knowledge_{i}: [ERROR DECODING]\n")
|
||
|
||
print(f"Decoding completed. Output saved to {output_path}")
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="Decode MiniMind model's knowledge database")
|
||
parser.add_argument("--model_path", type=str, default="out/pretrain_512.pth",
|
||
help="Path to the model checkpoint")
|
||
parser.add_argument("--output_path", type=str, default="out/knowledge_db.txt",
|
||
help="Path to save the decoded text file")
|
||
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
|
||
help="Device to load the model on")
|
||
|
||
args = parser.parse_args()
|
||
|
||
decode_dataset(args.model_path, args.output_path, args.device)
|