import json, shutil, sys, os, torch
from transformers import AutoConfig, AutoModelForCausalLM
src, out = sys.argv[1], sys.argv[2]
full = json.load(open(f"{src}/config.json"))
text = dict(full["text_config"])
text.update(hidden_size=64, intermediate_size=128, num_hidden_layers=4, num_attention_heads=2, num_key_value_heads=1,
            head_dim=32, linear_key_head_dim=16, linear_value_head_dim=16, linear_num_key_heads=2, linear_num_value_heads=2,
            layer_types=["linear_attention"] * 3 + ["full_attention"], mtp_num_hidden_layers=0)
text["model_type"] = "qwen3_5_text"
cfg = AutoConfig.for_model(**text)
torch.manual_seed(0)
model = AutoModelForCausalLM.from_config(cfg)
print(type(model).__name__, sum(p.numel() for p in model.parameters()) / 1e6, "M params")
model.save_pretrained(out)
for f in ["tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt"]:
    shutil.copy(f"{src}/{f}", out)