File size: 4,712 Bytes
4c383ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """Extract text weights, add six tokens and initialize a decision head. No training."""
import argparse
import hashlib
import json
import shutil
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import torch
from safetensors import safe_open
from transformers import AutoTokenizer
from transformers.dynamic_module_utils import get_class_from_dynamic_module
_REPO = str(Path(__file__).resolve().parents[1])
MugiDecisionConfig = get_class_from_dynamic_module('configuration_mugi.MugiDecisionConfig', _REPO, local_files_only=True)
MugiDecisionModel = get_class_from_dynamic_module('modeling_mugi.MugiDecisionModel', _REPO, local_files_only=True)
from mugi_decision.encoding import TOKENS
from mugi_decision.runtime import save_json
def main():
p = argparse.ArgumentParser()
p.add_argument('--source', type=Path, default=Path('Qwen3.5-0.8B'))
p.add_argument('--output', type=Path, default=Path('models/mugi-decision-0.8b-base'))
p.add_argument('--no-enum', action='store_true')
p.add_argument('--seed', type=int, default=42)
args = p.parse_args()
if args.output.exists() and any(args.output.iterdir()):
raise SystemExit('Output is not empty; select a new path.')
config = json.loads((args.source / 'config.json').read_text(encoding='utf-8'))
if config['model_type'] != 'qwen3_5':
raise ValueError('Expected an original Qwen3.5 multimodal checkpoint')
torch.manual_seed(args.seed)
tokenizer = AutoTokenizer.from_pretrained(args.source, local_files_only=True)
tokenizer.add_special_tokens({'extra_special_tokens': TOKENS}, replace_extra_special_tokens=False)
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = 'right'
token_ids = tokenizer.convert_tokens_to_ids(TOKENS)
cfg = MugiDecisionConfig(text_config=config['text_config'], include_enum=not args.no_enum,
decision_token_ids=token_ids, pad_token_id=tokenizer.pad_token_id)
cfg.text_config.pad_token_id = tokenizer.pad_token_id
cfg._attn_implementation = 'sdpa'
previous_dtype = torch.get_default_dtype()
try:
torch.set_default_dtype(torch.bfloat16)
model = MugiDecisionModel(cfg)
finally:
torch.set_default_dtype(previous_dtype)
weights = {}
prefix = 'model.language_model.'
index = args.source / 'model.safetensors.index.json'
filenames = sorted(set(json.loads(index.read_text())['weight_map'].values())) if index.exists() else [x.name for x in args.source.glob('*.safetensors')]
if not filenames: raise FileNotFoundError('No safetensors weights found')
for filename in filenames:
with safe_open(args.source / filename, framework='pt', device='cpu') as handle:
for key in handle.keys():
if key.startswith(prefix):
weights[key[len(prefix):]] = handle.get_tensor(key)
model.backbone.load_state_dict(weights, strict=True, assign=True)
old_size = model.get_input_embeddings().num_embeddings
size = max(old_size, len(tokenizer))
# Preserve existing vocabulary rows; initialize new markers in unused padded rows.
model.resize_token_embeddings(size, mean_resizing=False)
cfg.text_config.vocab_size = size
model.backbone.config.vocab_size = size
with torch.no_grad():
for token_id in token_ids:
torch.nn.init.normal_(model.get_input_embeddings().weight[token_id], std=cfg.initializer_range)
MugiDecisionConfig.register_for_auto_class()
MugiDecisionModel.register_for_auto_class('AutoModel')
args.output.mkdir(parents=True, exist_ok=True)
model.save_pretrained(args.output, safe_serialization=True, max_shard_size='4GB')
tokenizer.save_pretrained(args.output)
if (args.source / 'LICENSE').exists(): shutil.copyfile(args.source / 'LICENSE', args.output / 'LICENSE')
metadata = dict(source=str(args.source.resolve()), source_config_sha256=hashlib.sha256((args.source/'config.json').read_bytes()).hexdigest(), source_weight_files=filenames,
seed=args.seed, parameters=sum(x.numel() for x in model.parameters()),
original_vocab_size=old_size, vocab_size=size, special_tokens=dict(zip(TOKENS, token_ids)),
include_enum=cfg.include_enum, state='initialized_not_finetuned',
note='Text backbone copied strictly. Vision encoder and vocabulary output projection omitted. Decision head and marker embeddings require training.')
save_json(args.output / 'conversion.json', metadata)
print(json.dumps(metadata, ensure_ascii=False, indent=2))
if __name__ == '__main__': main()
|