Update num_classes in JTP2Processor to 18166 for alignment with state_dict

caption/jtp2.py

@@ -94,7 +94,7 @@ class JTP2Processor(BatchProcessor[Path, None]):
         self.model = timm.create_model(
             "vit_so400m_patch14_siglip_384.webli",
             pretrained=False,
-            num_classes=9083
+            num_classes=18166  # Align with state_dict
         )
         
         # Replace the model's head with the custom GatedHead

verify_tags.py

@@ -0,0 +1,5 @@
+import json
+
+with open("/home/kade/source/repos/JTP2/tags.json", "r", encoding="utf-8") as f:
+    tags = json.load(f)
+print(f"Number of tags: {len(tags)}")