import json
from pathlib import Path

# Paths (adjust if your files are elsewhere)
config_path = Path("config.json")
vocab_path = Path("vocab.json")

# Load files
with open(config_path, "r", encoding="utf-8") as f:
    config = json.load(f)
with open(vocab_path, "r", encoding="utf-8") as f:
    vocab = json.load(f)

# vocab.json is usually {"AA": 1, "AE": 2, ...}
# we need to invert it to { "1": "AA", "2": "AE", ... }
id2label = {str(idx): token for token, idx in vocab.items()}

# Insert into config.json
config["id2label"] = id2label
config["label2id"] = vocab   # optional, nice for completeness

# Save back
with open(config_path, "w", encoding="utf-8") as f:
    json.dump(config, f, indent=2, ensure_ascii=False)

print(f"Added id2label with {len(id2label)} entries to {config_path}")