feat: convert .pt checkpoints to HuggingFace-style state_dict + config

Convert pickle-based full model .pt files to safetensors + config.json
structure for PyTorch version compatibility and safer loading.

Datasets: MSL, SMAP, SWaT, WADI

MSL/config.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "AnomalyBERT",
+  "input_d_data": 55,
+  "output_d_data": 1,
+  "patch_size": 2,
+  "d_embed": 512,
+  "hidden_dim_rate": 4.0,
+  "max_seq_len": 512,
+  "positional_encoding": null,
+  "relative_position_embedding": true,
+  "transformer_n_layer": 6,
+  "transformer_n_head": 8,
+  "dropout": 0.1
+}

MSL/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:732a51712cd2a1402a0c18cba2b895fa5daf437229c561801bc875b50dc2ade4
+size 80314744

SMAP/config.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "AnomalyBERT",
+  "input_d_data": 25,
+  "output_d_data": 1,
+  "patch_size": 4,
+  "d_embed": 512,
+  "hidden_dim_rate": 4.0,
+  "max_seq_len": 512,
+  "positional_encoding": null,
+  "relative_position_embedding": true,
+  "transformer_n_layer": 6,
+  "transformer_n_head": 8,
+  "dropout": 0.1
+}

SMAP/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:139feda8feff168b1eeb2e4bb977e93bbe25bb46b7ebff31b9073dbd3e1b3a08
+size 80310656

SWaT/config.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "AnomalyBERT",
+  "input_d_data": 50,
+  "output_d_data": 1,
+  "patch_size": 14,
+  "d_embed": 512,
+  "hidden_dim_rate": 4.0,
+  "max_seq_len": 512,
+  "positional_encoding": null,
+  "relative_position_embedding": true,
+  "transformer_n_layer": 6,
+  "transformer_n_head": 8,
+  "dropout": 0.1
+}

SWaT/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b80bbf2b36004c9061957b408765645262ddbf129edab0b4baab148efcb78b09
+size 81621424

WADI/config.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "AnomalyBERT",
+  "input_d_data": 122,
+  "output_d_data": 1,
+  "patch_size": 8,
+  "d_embed": 512,
+  "hidden_dim_rate": 4.0,
+  "max_seq_len": 512,
+  "positional_encoding": null,
+  "relative_position_embedding": true,
+  "transformer_n_layer": 6,
+  "transformer_n_head": 8,
+  "dropout": 0.1
+}

WADI/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83eca0fd163b5bf7c9bc1a564a98dd0fb21e4b838687bb4fa724f488dd900201
+size 82137496

convert_to_hf.py

@@ -0,0 +1,122 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = ["torch", "timm", "safetensors"]
+# ///
+"""Convert AnomalyBERT .pt checkpoint files to HuggingFace-style state_dict + config structure.
+
+Each .pt file is converted to a directory containing:
+  - config.json: model hyperparameters
+  - model.safetensors: state_dict in safetensors format
+"""
+
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, 'original')
+
+import torch
+from safetensors.torch import save_file
+
+
+def extract_config(model: torch.nn.Module) -> dict:
+    """Extract model hyperparameters from a loaded AnomalyTransformer."""
+    patch_size = model.patch_size
+    max_seq_len = model.max_seq_len
+    d_embed = model.linear_embedding.weight.shape[0]
+    input_d_data = model.linear_embedding.weight.shape[1] // patch_size
+
+    # output_d_data from mlp_layers[-1].out_features / patch_size
+    mlp_last = model.mlp_layers[-1]
+    output_d_data = mlp_last.out_features // patch_size
+
+    # Count transformer layers
+    n_layer = len(model.transformer_encoder.encoder_layers)
+
+    # Get n_head from first attention layer
+    first_attn = model.transformer_encoder.encoder_layers[0].attention_layer
+    n_head = first_attn.n_head
+
+    # Get hidden_dim from first feed forward layer
+    first_ff = model.transformer_encoder.encoder_layers[0].feed_forward_layer
+    hidden_dim = first_ff.first_fc_layer.out_features
+    hidden_dim_rate = hidden_dim / d_embed
+
+    # Detect positional encoding type
+    has_pe = model.transformer_encoder.positional_encoding
+    if has_pe:
+        pe_layer = model.transformer_encoder.positional_encoding_layer
+        pe_type = type(pe_layer).__name__
+        if 'Sinusoidal' in pe_type:
+            positional_encoding = 'Sinusoidal'
+        elif 'Absolute' in pe_type:
+            positional_encoding = 'Absolute'
+        else:
+            positional_encoding = pe_type
+    else:
+        positional_encoding = None
+
+    # Detect relative position embedding
+    relative_position_embedding = first_attn.relative_position_embedding
+
+    # Dropout rate from encoder layer
+    dropout = model.transformer_encoder.encoder_layers[0].dropout_layer.p
+
+    return {
+        'model_type': 'AnomalyBERT',
+        'input_d_data': input_d_data,
+        'output_d_data': output_d_data,
+        'patch_size': patch_size,
+        'd_embed': d_embed,
+        'hidden_dim_rate': hidden_dim_rate,
+        'max_seq_len': max_seq_len,
+        'positional_encoding': positional_encoding,
+        'relative_position_embedding': relative_position_embedding,
+        'transformer_n_layer': n_layer,
+        'transformer_n_head': n_head,
+        'dropout': dropout,
+    }
+
+
+def convert_checkpoint(pt_path: Path, output_dir: Path) -> None:
+    """Convert a single .pt checkpoint to state_dict + config."""
+    print(f'Converting {pt_path.name}...')
+
+    model = torch.load(pt_path, map_location='cpu', weights_only=False)
+    config = extract_config(model)
+    state_dict = model.state_dict()
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Save config
+    config_path = output_dir / 'config.json'
+    with open(config_path, 'w') as f:
+        json.dump(config, f, indent=2)
+    print(f'  config.json: {json.dumps(config, indent=2)}')
+
+    # Save state_dict as safetensors
+    safetensors_path = output_dir / 'model.safetensors'
+    save_file(state_dict, str(safetensors_path))
+    size_mb = safetensors_path.stat().st_size / 1024 / 1024
+    print(f'  model.safetensors: {size_mb:.1f} MB ({len(state_dict)} tensors)')
+
+    print(f'  -> {output_dir}/')
+
+
+def main() -> None:
+    """Convert all *_parameters.pt files."""
+    pt_files = sorted(Path('.').glob('*_parameters.pt'))
+    if not pt_files:
+        print('No *_parameters.pt files found.')
+        return
+
+    for pt_path in pt_files:
+        dataset_name = pt_path.stem.replace('_parameters', '')
+        output_dir = Path(dataset_name)
+        convert_checkpoint(pt_path, output_dir)
+
+    print(f'\nDone. Converted {len(pt_files)} checkpoints.')
+
+
+if __name__ == '__main__':
+    main()

inspect_pt.py

@@ -0,0 +1,47 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = ["torch", "timm"]
+# ///
+"""Inspect the structure of .pt checkpoint files."""
+
+import sys
+
+sys.path.insert(0, 'original')
+
+import torch
+from pathlib import Path
+
+
+def inspect_checkpoint(path: Path) -> None:
+    """Print the structure of a checkpoint file."""
+    print(f'\n{"=" * 60}')
+    print(f'File: {path.name} ({path.stat().st_size / 1024 / 1024:.1f} MB)')
+    print('=' * 60)
+
+    data = torch.load(path, map_location='cpu', weights_only=False)
+    print(f'Top-level type: {type(data).__name__}')
+
+    if isinstance(data, dict):
+        print(f'Keys: {list(data.keys())[:10]}')
+        for key, val in list(data.items())[:5]:
+            if isinstance(val, torch.Tensor):
+                print(f'  {key}: Tensor {val.shape} {val.dtype}')
+            else:
+                print(f'  {key}: {type(val).__name__}')
+    elif hasattr(data, 'state_dict'):
+        print('This is a full model object.')
+        sd = data.state_dict()
+        print(f'state_dict keys ({len(sd)}):')
+        for k, v in sd.items():
+            print(f'  {k}: {v.shape} {v.dtype}')
+        # Print model attributes
+        for attr in ['max_seq_len', 'patch_size', 'data_seq_len']:
+            if hasattr(data, attr):
+                print(f'  model.{attr} = {getattr(data, attr)}')
+    else:
+        print(f'Unexpected type: {type(data)}')
+
+
+if __name__ == '__main__':
+    for pt_file in sorted(Path('.').glob('*_parameters.pt')):
+        inspect_checkpoint(pt_file)

verify_conversion.py

@@ -0,0 +1,85 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = ["torch", "timm", "safetensors"]
+# ///
+"""Verify converted safetensors match original .pt checkpoints."""
+
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, 'original')
+
+import torch
+from safetensors.torch import load_file
+
+from models.anomaly_transformer import get_anomaly_transformer
+
+
+def verify(dataset: str) -> bool:
+    """Verify a single converted checkpoint."""
+    pt_path = Path(f'{dataset}_parameters.pt')
+    config_path = Path(dataset) / 'config.json'
+    safetensors_path = Path(dataset) / 'model.safetensors'
+
+    # Load original
+    original = torch.load(pt_path, map_location='cpu', weights_only=False)
+    original_sd = original.state_dict()
+
+    # Load config and rebuild model
+    with open(config_path) as f:
+        config = json.load(f)
+
+    model = get_anomaly_transformer(
+        input_d_data=config['input_d_data'],
+        output_d_data=config['output_d_data'],
+        patch_size=config['patch_size'],
+        d_embed=config['d_embed'],
+        hidden_dim_rate=config['hidden_dim_rate'],
+        max_seq_len=config['max_seq_len'],
+        positional_encoding=config['positional_encoding'],
+        relative_position_embedding=config['relative_position_embedding'],
+        transformer_n_layer=config['transformer_n_layer'],
+        transformer_n_head=config['transformer_n_head'],
+        dropout=config['dropout'],
+    )
+
+    # Load safetensors weights
+    saved_sd = load_file(str(safetensors_path))
+    model.load_state_dict(saved_sd)
+    loaded_sd = model.state_dict()
+
+    # Compare
+    ok = True
+    for key in original_sd:
+        if key not in loaded_sd:
+            print(f'  MISSING: {key}')
+            ok = False
+            continue
+        if not torch.equal(original_sd[key], loaded_sd[key]):
+            diff = (original_sd[key] - loaded_sd[key]).abs().max().item()
+            print(f'  MISMATCH: {key} (max diff={diff})')
+            ok = False
+
+    extra = set(loaded_sd.keys()) - set(original_sd.keys())
+    if extra:
+        print(f'  EXTRA keys: {extra}')
+        ok = False
+
+    status = 'OK' if ok else 'FAIL'
+    print(f'{dataset}: {status}')
+    return ok
+
+
+def main() -> None:
+    """Verify all converted checkpoints."""
+    datasets = ['MSL', 'SMAP', 'SWaT', 'WADI']
+    results = {d: verify(d) for d in datasets}
+    all_ok = all(results.values())
+    print(f'\nAll passed: {all_ok}')
+    if not all_ok:
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    main()