| """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)) |
| |
| 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() |
|
|
