Support architecture override fine-tuning

anifilebert/train.py

@@ -1264,6 +1264,45 @@ def build_vocab_from_data(data: List[Dict], tokenizer: AnimeTokenizer, vocab_pat
         json.dump(tokenizer.get_vocab(), f, ensure_ascii=False, indent=2)
 
 
+def architecture_matches_config(model, config: Config) -> bool:
+    model_config = model.config
+    return (
+        int(getattr(model_config, "hidden_size", -1)) == int(config.hidden_size)
+        and int(getattr(model_config, "num_hidden_layers", -1)) == int(config.num_hidden_layers)
+        and int(getattr(model_config, "num_attention_heads", -1)) == int(config.num_attention_heads)
+        and int(getattr(model_config, "intermediate_size", -1)) == int(config.intermediate_size)
+        and int(getattr(model_config, "max_position_embeddings", -1)) >= int(config.max_position_embeddings)
+    )
+
+
+def rebuild_model_with_architecture_overrides(source_model, config: Config, model_head: str):
+    """Create a target architecture and copy same-shaped tensors from source."""
+    target_model = create_model(config, model_head=model_head)
+    source_state = source_model.state_dict()
+    target_state = target_model.state_dict()
+    copied_tensors = 0
+    copied_parameters = 0
+    skipped = []
+
+    with torch.no_grad():
+        for name, target_tensor in target_state.items():
+            source_tensor = source_state.get(name)
+            if source_tensor is None or tuple(source_tensor.shape) != tuple(target_tensor.shape):
+                skipped.append(name)
+                continue
+            target_tensor.copy_(source_tensor.to(device=target_tensor.device, dtype=target_tensor.dtype))
+            copied_tensors += 1
+            copied_parameters += target_tensor.numel()
+
+    target_model.load_state_dict(target_state)
+    return target_model, {
+        "copied_tensors": copied_tensors,
+        "copied_parameters": copied_parameters,
+        "skipped_tensors": len(skipped),
+        "skipped_examples": skipped[:10],
+    }
+
+
 def main():
     args = parse_args()
     config = Config()
@@ -1434,6 +1473,25 @@ def main():
         model.config.id2label = config.id2label
         model.config.label2id = config.label2id
         model.config.label_schema_version = config.label_schema_version
+        if not architecture_matches_config(model, config):
+            print(
+                "  Rebuilding model for architecture overrides: "
+                f"layers={config.num_hidden_layers}, heads={config.num_attention_heads}, "
+                f"hidden={config.hidden_size}, intermediate={config.intermediate_size}"
+            )
+            model, architecture_copy = rebuild_model_with_architecture_overrides(
+                source_model=model,
+                config=config,
+                model_head=selected_model_head,
+            )
+            print(
+                "  Copied compatible tensors into override architecture: "
+                f"{architecture_copy['copied_tensors']} tensors, "
+                f"{architecture_copy['copied_parameters']:,} params; "
+                f"skipped {architecture_copy['skipped_tensors']} tensors"
+            )
+            if architecture_copy["skipped_examples"]:
+                print(f"  Skipped tensor examples: {architecture_copy['skipped_examples']}")
     else:
         print("Creating model...")
         selected_model_head = "linear" if args.model_head == "auto" else args.model_head