Spaces:
Sleeping
Sleeping
| """ | |
| 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__) | |
| class TokenImportance: | |
| token: str | |
| score: float | |
| class ModelPrediction: | |
| model_name: str | |
| label: str | |
| confidence: float | |
| scores: Dict[str, float] | |
| tokens: List[Dict] | |
| 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 | |