""" Ensemble classifier combining DistilBERT, RoBERTa, and XLNet with parallel execution and multiple voting strategies. """ import asyncio import time import logging from dataclasses import dataclass from typing import List, Dict, Optional from concurrent.futures import ThreadPoolExecutor import numpy as np from .inference import get_classifier logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @dataclass class TokenImportance: token: str score: float @dataclass class ModelPrediction: model_name: str label: str confidence: float scores: Dict[str, float] tokens: List[Dict] @dataclass class EnsembleResult: hard_voting_label: str hard_voting_confidence: float soft_voting_label: str soft_voting_confidence: float soft_voting_scores: Dict[str, float] weighted_voting_label: str weighted_voting_confidence: float weighted_voting_scores: Dict[str, float] uncertainty: float model_agreement: float individual_predictions: List[ModelPrediction] merged_explanation: List[Dict] execution_time_ms: float warnings: Optional[List[str]] = None class EnsembleClassifier: """Combines predictions from all three models using voting strategies.""" def __init__(self): self.model_names = ['distilbert', 'roberta', 'xlnet'] self.models = {name: get_classifier(name) for name in self.model_names} # Per-class F1 scores from training (more accurate than overall accuracy) self.class_weights = { 'distilbert': {'True': 0.889, 'Fake': 0.872, 'Satire': 0.998, 'Bias': 0.633}, 'roberta': {'True': 0.888, 'Fake': 0.879, 'Satire': 0.998, 'Bias': 0.615}, 'xlnet': {'True': 0.892, 'Fake': 0.876, 'Satire': 0.997, 'Bias': 0.638} } self.executor = ThreadPoolExecutor(max_workers=3) async def predict_ensemble(self, text: str, model_timeout: float = 60.0, total_timeout: float = 180.0) -> EnsembleResult: start_time = time.time() warnings = [] # Check if running on CPU - if so, run sequentially instead of parallel import torch use_sequential = not torch.cuda.is_available() if use_sequential: # Sequential execution on CPU (more efficient) logger.info("[ensemble] Running models sequentially on CPU") results = [] for name in self.model_names: try: result = self.models[name].predict(text) results.append(result) except Exception as e: logger.error(f"[ensemble] {name} prediction failed: {e}") results.append(None) else: # Parallel execution on GPU logger.info("[ensemble] Running models in parallel on GPU") loop = asyncio.get_event_loop() tasks = [ loop.run_in_executor( self.executor, self._predict_with_timeout, name, text, model_timeout) for name in self.model_names ] try: results = await asyncio.wait_for( asyncio.gather(*tasks, return_exceptions=True), timeout=total_timeout ) except asyncio.TimeoutError: warnings.append("Ensemble prediction exceeded total timeout") raise valid_predictions = [] for name, result in zip(self.model_names, results): if isinstance(result, Exception): warnings.append(f"{name} failed: {str(result)}") elif result is None: warnings.append(f"{name} returned no result") else: valid_predictions.append(ModelPrediction( model_name=name, label=result['label'], confidence=result['confidence'], scores=result['scores'], tokens=result['tokens'] )) if not valid_predictions: raise RuntimeError("All models failed to process the request") hard_label, hard_conf = self.hard_voting(valid_predictions) soft_scores = self.soft_voting(valid_predictions) soft_label = max(soft_scores.items(), key=lambda x: x[1])[0] soft_conf = soft_scores[soft_label] weighted_scores = self.weighted_voting(valid_predictions) weighted_label = max(weighted_scores.items(), key=lambda x: x[1])[0] weighted_conf = weighted_scores[weighted_label] # Calculate uncertainty (1 - max confidence across all voting methods) max_confidence = max(hard_conf, soft_conf, weighted_conf) uncertainty = round(1.0 - max_confidence, 4) # Calculate model agreement (what % of models predicted the most common label) label_counts = {} for pred in valid_predictions: label_counts[pred.label] = label_counts.get(pred.label, 0) + 1 most_common_count = max(label_counts.values()) model_agreement = round(most_common_count / len(valid_predictions), 4) merged_tokens = self._merge_explanations( valid_predictions, weighted_label) execution_time = (time.time() - start_time) * 1000 logger.info( f"Ensemble completed in {execution_time:.2f}ms with {len(valid_predictions)}/{len(self.model_names)} models") if warnings: logger.warning(f"Ensemble warnings: {warnings}") return EnsembleResult( hard_voting_label=hard_label, hard_voting_confidence=hard_conf, soft_voting_label=soft_label, soft_voting_confidence=soft_conf, soft_voting_scores=soft_scores, weighted_voting_label=weighted_label, weighted_voting_confidence=weighted_conf, weighted_voting_scores=weighted_scores, uncertainty=uncertainty, model_agreement=model_agreement, individual_predictions=valid_predictions, merged_explanation=merged_tokens, execution_time_ms=round(execution_time, 2), warnings=warnings if warnings else None ) def _predict_with_timeout(self, model_name: str, text: str, timeout: float) -> Optional[Dict]: """ Run model prediction with timeout enforcement. Uses concurrent.futures to enforce per-model timeout. """ try: from concurrent.futures import TimeoutError as FuturesTimeoutError # Submit to executor with timeout future = self.executor.submit( self.models[model_name].predict, text) return future.result(timeout=timeout) except FuturesTimeoutError: logger.error( f"[ensemble] {model_name} prediction timed out after {timeout}s") return None except Exception as e: logger.error(f"[ensemble] {model_name} prediction failed: {e}") return None def hard_voting(self, predictions: List[ModelPrediction]) -> tuple[str, float]: """ Hard voting: select label with most votes. In case of tie, select label with highest average confidence among tied labels. """ votes = {} for pred in predictions: votes[pred.label] = votes.get(pred.label, 0) + 1 # Find maximum vote count max_votes = max(votes.values()) # Get all labels with max votes (handles ties) tied_labels = [label for label, count in votes.items() if count == max_votes] # If tie, break by highest average confidence if len(tied_labels) > 1: label_confidences = {} for label in tied_labels: confidences = [ p.confidence for p in predictions if p.label == label] label_confidences[label] = sum(confidences) / len(confidences) winning_label = max(label_confidences.items(), key=lambda x: x[1])[0] else: winning_label = tied_labels[0] # Average confidence of models that predicted the winning label confidences = [ p.confidence for p in predictions if p.label == winning_label] return winning_label, round(sum(confidences) / len(confidences), 4) def soft_voting(self, predictions: List[ModelPrediction]) -> Dict[str, float]: """ Confidence-weighted soft voting: averages probabilities weighted by model confidence. Models that are more confident have more influence. """ all_labels = set( label for pred in predictions for label in pred.scores) weighted_scores = {} total_confidence = sum(p.confidence for p in predictions) for label in all_labels: weighted_sum = sum( p.scores.get(label, 0.0) * p.confidence for p in predictions ) weighted_scores[label] = round(weighted_sum / total_confidence, 4) return weighted_scores def weighted_voting(self, predictions: List[ModelPrediction]) -> Dict[str, float]: """ Per-class F1-weighted voting: uses class-specific F1 scores as weights. More accurate than using overall model accuracy. """ all_labels = set( label for pred in predictions for label in pred.scores) weighted_scores = {} for label in all_labels: weighted_sum = 0.0 total_weight = 0.0 for pred in predictions: prob = pred.scores.get(label, 0.0) f1_weight = self.class_weights[pred.model_name][label] weighted_sum += prob * f1_weight total_weight += f1_weight weighted_scores[label] = round(weighted_sum / total_weight, 4) return weighted_scores def _merge_explanations(self, predictions: List[ModelPrediction], predicted_label: str) -> List[Dict]: """ Merge token importance scores across models, weighted by class-specific F1 scores. Tokens from models with better performance on the predicted class get more weight. """ token_scores: Dict[str, float] = {} token_counts: Dict[str, int] = {} for pred in predictions: # Weight tokens by the model's F1 score for the predicted class f1_weight = self.class_weights[pred.model_name][predicted_label] for td in pred.tokens: token = td['token'] weighted_score = td['score'] * f1_weight token_scores[token] = token_scores.get( token, 0.0) + weighted_score token_counts[token] = token_counts.get(token, 0) + 1 merged = [ {'token': t, 'score': round(token_scores[t] / token_counts[t], 4)} for t in token_scores ] return sorted(merged, key=lambda x: x['score'], reverse=True)[:10] _ensemble_classifier: Optional[EnsembleClassifier] = None def get_ensemble_classifier() -> EnsembleClassifier: global _ensemble_classifier if _ensemble_classifier is None: _ensemble_classifier = EnsembleClassifier() return _ensemble_classifier