import json
from safetensors.torch import load_file as load_safetensors
# from your_code.gptmoe import GPTMoE  # ← あなたの実装を import

with open("config.json", "r", encoding="utf-8") as f:
    cfg = json.load(f)
moe = cfg["moe"]
model = GPTMoE(
    cfg["vocab_size"], cfg["d_model"], cfg["n_heads"], cfg["n_layers"], cfg["ffn_mult"],
    dict(
        num_experts=moe["num_experts"], k=moe["top_k"],
        capacity_factor=moe["capacity_factor"], eval_capacity_factor=moe["eval_capacity_factor"],
        min_capacity=0, noisy_gate_policy=moe["noisy_gate_policy"], use_residual=moe["use_residual"],
    ),
).eval()
state = load_safetensors("model.safetensors")
model.load_state_dict(state, strict=True)

import sentencepiece as spm
sp = spm.SentencePieceProcessor(model_file="tokenizer/ja_unigram32k_v15m.model")
ids = sp.encode("こんにちは", out_type=int)