Update model_loader.py

model_loader.py

@@ -225,75 +225,7 @@ class SemanticFeatureExtractor:
 # ============================================================
 
 if TORCH_AVAILABLE:
-    class EOUClassifier(nn.Module):
-        """AutoModel with auxiliary features for End-of-Utterance detection"""
-
-        def __init__(self, model_path: str, use_aux: bool = True, num_aux_features: int = 15, dropout: float = 0.1):  
-            super().__init__()
-            self.use_aux = use_aux
-
-            # Load the configuration purely from local files
-            from transformers import AutoConfig, AutoModel
-            config = AutoConfig.from_pretrained(model_path, local_files_only=True)
-            
-            # Initialize model architecture WITHOUT downloading base weights
-            self.base_model = AutoModel.from_config(config)
-            
-            # DistilBert uses 'dim', others use 'hidden_size'
-            hidden_size = getattr(config, 'hidden_size', getattr(config, 'dim', 768))
-
-            self.pooler_dropout = nn.Dropout(dropout)
-
-            if self.use_aux:
-                self.aux_projection = nn.Sequential(
-                    nn.Linear(num_aux_features, 32),
-                    nn.GELU(),
-                    nn.Dropout(dropout),
-                )
-                classifier_input_size = hidden_size + 32
-            else:
-                classifier_input_size = hidden_size
-
-            self.classifier = nn.Sequential(
-                nn.Linear(classifier_input_size, 256),
-                nn.GELU(),
-                nn.LayerNorm(256),
-                nn.Dropout(dropout),
-                nn.Linear(256, 64),
-                nn.GELU(),
-                nn.Dropout(dropout),
-                nn.Linear(64, 2),
-            )
-
-        def forward(self, input_ids, attention_mask, token_type_ids=None,
-                    aux_features=None, labels=None):
-
-            # DistilBert doesn't accept token_type_ids
-            model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask}
-            if token_type_ids is not None and "token_type_ids" in self.base_model.forward.__code__.co_varnames:
-                 model_inputs["token_type_ids"] = token_type_ids
-                 
-            outputs = self.base_model(**model_inputs)
-
-            # Get the CLS token representation (first token)
-            cls_output = outputs.last_hidden_state[:, 0, :]
-            cls_output = self.pooler_dropout(cls_output)
-
-            if self.use_aux and aux_features is not None:
-                aux_projected = self.aux_projection(aux_features)
-                combined = torch.cat([cls_output, aux_projected], dim=-1)
-            else:
-                combined = cls_output
-
-            logits = self.classifier(combined)
-
-            loss = None
-            if labels is not None:
-                # Default smoothing if none configured
-                loss_fn = nn.CrossEntropyLoss(label_smoothing=0.05)
-                loss = loss_fn(logits, labels)
-
-            return {'loss': loss, 'logits': logits}
+    from transformers import AutoModelForSequenceClassification
 
 
 
@@ -376,32 +308,13 @@ class EOUModelEngine:
 
                     model_config = Config()
                     model_config.model_name = self.eou_config.get(
-                        'model_name', 'microsoft/deberta-v3-base'
-                    )
-                    model_config.use_aux_features = self.eou_config.get(
-                        'use_aux_features', True
-                    )
-                    num_aux = self.eou_config.get('num_aux_features', 15)
-
                     def _load_pytorch():
-                        model = EOUClassifier(
-                            model_path=model_dir, 
-                            use_aux=self.eou_config.get('use_aux_features', True),
-                            num_aux_features=num_aux
+                        # The model is a standard HF classification model (like DistilBertForSequenceClassification)
+                        # This natively handles config.json AND strictly loads your model.safetensors weights!
+                        model = AutoModelForSequenceClassification.from_pretrained(
+                            model_dir, 
+                            local_files_only=True
                         )
-                        
-                        # Try to find weights
-                        for alt in ['model.safetensors', 'pytorch_model.bin', 'pytorch_model_full.pt']:
-                            alt_path = os.path.join(model_dir, alt)
-                            if os.path.exists(alt_path):
-                                if alt.endswith('.safetensors'):
-                                    state_dict = load_safetensors(alt_path, device=str(self.device))
-                                else:
-                                    state_dict = torch.load(alt_path, map_location=self.device, weights_only=True)
-                                break
-                        else:
-                            raise FileNotFoundError(f"No model weights found in {model_dir}")
-                        model.load_state_dict(state_dict, strict=False)
                         model.to(self.device)
                         model.eval()
                         return model
@@ -564,19 +477,17 @@ class EOUModelEngine:
         if token_type_ids is not None:
             token_type_ids = token_type_ids.to(self.device)
 
-        aux_features = torch.tensor(
-            [self.feature_extractor.extract(clean_text)], dtype=torch.float
-        ).to(self.device)
-
         with torch.no_grad():
-            outputs = self.torch_model(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                token_type_ids=token_type_ids,
-                aux_features=aux_features,
-            )
+            model_inputs = {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask
+            }
+            if token_type_ids is not None and "token_type_ids" in self.torch_model.forward.__code__.co_varnames:
+                model_inputs["token_type_ids"] = token_type_ids
+
+            outputs = self.torch_model(**model_inputs)
 
-        probs = torch.softmax(outputs['logits'], dim=-1)[0].cpu().numpy()
+        probs = torch.softmax(outputs.logits, dim=-1)[0].cpu().numpy()
         complete_prob = float(probs[1])
         incomplete_prob = float(probs[0])
         is_complete = complete_prob >= self.threshold