File size: 1,517 Bytes
0ab2566 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """
CODI + Direct Soft Matching (Ablation: No Prototypes)
Inherits from CODI_CT and replaces PMPSLatentSupervisor with
DirectMatchSupervisor. The forward() method is fully reused — only
the pmps_module is swapped at init time.
Usage:
from src.model_direct_match import CODI_DirectMatch
model = CODI_DirectMatch(model_args, training_args, lora_config)
"""
from .model import CODI_CT
from .direct_match_loss import DirectMatchSupervisor
class CODI_DirectMatch(CODI_CT):
"""CODI with direct soft matching (no prototypes).
Ablation variant: replaces the prototype-mediated PMPS module with
DirectMatchSupervisor, which computes soft many-to-many matching
directly in the projected embedding space via cosine similarity.
All other components (base model, LoRA, projector, PSA, diversity
loss, distillation loss, CE losses) remain identical.
"""
def __init__(self, model_args, training_args, lora_config):
super().__init__(model_args, training_args, lora_config)
if self.pmps_loss_factor > 0:
del self.pmps_module
self.pmps_module = DirectMatchSupervisor(
hidden_dim=self.dim,
feat_dim=training_args.pmps_feat_dim,
hidden_mlp=training_args.pmps_hidden_mlp,
match_temperature=training_args.pmps_match_temperature,
use_projector=training_args.pmps_use_projector,
position_bias_strength=training_args.pmps_position_bias_max,
)
|