Alternate simpler fix for jailbreak model loading issue

[srampal]

Strip _orig_mod when saving instead of during retrieval as done in prior fix attempt
Alternate fix for #54

[srampal]

This fix also includes a self-test that ensures the correct model parameters are saved.

This was verified by fine tuning a model and verifying that the model parameters saved are the same as the original ones without the torch compile prefixes as seen in the output of training below.

{'train_runtime': 131505.762, 'train_samples_per_second': 81.284, 'train_steps_per_second': 3.695, 'train_loss': 0.0005007739498931483, 'epoch': 10.0}
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 485870/485870 [36:31:45<00:00, 3.69it/s]
INFO:main:Training completed successfully!
INFO:main:Detected torch.compile wrapper, extracting original model for saving...
INFO:main:Parameter verification: True
INFO:main:Model and tokenizer saved to jailbreak_classifier_modernbert-base_model without compile prefixes
INFO:main:Evaluating on validation set...
W0824 14:11:52.817000 82749 torch/fx/experimental/symbolic_shapes.py:6823] [2/3] _maybe_guard_rel() was called on non-relation expression Eq(s53, 1) | Eq(s70, s53)
INFO:main:Validation accuracy: 0.9998
INFO:main:Evaluating on test set...
INFO:main:Test accuracy: 0.9998

==================================================
Jailbreak Classification Fine-tuning Completed!

Validation Accuracy: 0.9998
Test Accuracy: 0.9998

[srampal]

Someone with admin privileges needs to address the workflow permissions error (Owner Notification ...), I believe.

[Xunzhuo]

Author

/cc @rootfs