# DEPENDENCIES from typing import Dict from typing import List from typing import Tuple from loguru import logger from typing import Optional from config.enums import Domain from config.schemas import DomainPrediction from models.model_manager import get_model_manager from config.constants import domain_classification_params from config.threshold_config import interpolate_thresholds from config.threshold_config import get_threshold_for_domain class DomainClassifier: """ Classifies text into domains using zero-shot classification """ # Use constants from config - map string keys to Domain enum DOMAIN_LABELS = {Domain.ACADEMIC : domain_classification_params.DOMAIN_LABELS["academic"], Domain.CREATIVE : domain_classification_params.DOMAIN_LABELS["creative"], Domain.AI_ML : domain_classification_params.DOMAIN_LABELS["ai_ml"], Domain.SOFTWARE_DEV : domain_classification_params.DOMAIN_LABELS["software_dev"], Domain.TECHNICAL_DOC : domain_classification_params.DOMAIN_LABELS["technical_doc"], Domain.ENGINEERING : domain_classification_params.DOMAIN_LABELS["engineering"], Domain.SCIENCE : domain_classification_params.DOMAIN_LABELS["science"], Domain.BUSINESS : domain_classification_params.DOMAIN_LABELS["business"], Domain.JOURNALISM : domain_classification_params.DOMAIN_LABELS["journalism"], Domain.SOCIAL_MEDIA : domain_classification_params.DOMAIN_LABELS["social_media"], Domain.BLOG_PERSONAL : domain_classification_params.DOMAIN_LABELS["blog_personal"], Domain.LEGAL : domain_classification_params.DOMAIN_LABELS["legal"], Domain.MEDICAL : domain_classification_params.DOMAIN_LABELS["medical"], Domain.MARKETING : domain_classification_params.DOMAIN_LABELS["marketing"], Domain.TUTORIAL : domain_classification_params.DOMAIN_LABELS["tutorial"], Domain.GENERAL : domain_classification_params.DOMAIN_LABELS["general"], } def __init__(self): self.model_manager = get_model_manager() self.primary_classifier = None self.fallback_classifier = None self.is_initialized = False def initialize(self) -> bool: """ Initialize the domain classifier with zero-shot models """ try: logger.info("Initializing domain classifier...") # Load primary domain classifier (zero-shot) self.primary_classifier = self.model_manager.load_model(model_name = "content_domain_classifier") # Load fallback classifier try: self.fallback_classifier = self.model_manager.load_model(model_name = "domain_classifier_fallback") logger.info("Fallback classifier loaded successfully") except Exception as e: logger.warning(f"Could not load fallback classifier: {repr(e)}") self.fallback_classifier = None self.is_initialized = True logger.success("Domain classifier initialized successfully") return True except Exception as e: logger.error(f"Failed to initialize domain classifier: {repr(e)}") return False def classify(self, text: str, top_k: int = domain_classification_params.TOP_K_DOMAINS, min_confidence: float = domain_classification_params.MIN_CONFIDENCE_THRESHOLD) -> DomainPrediction: """ Classify text into domain using zero-shot classification Arguments: ---------- text { str } : Input text top_k { int } : Number of top domains to consider min_confidence { float } : Minimum confidence threshold Returns: -------- { DomainPrediction } : DomainPrediction object """ if not self.is_initialized: logger.warning("Domain classifier not initialized, initializing now...") if not self.initialize(): return self._get_default_prediction() try: # First try with primary classifier primary_result = self._classify_with_model(text = text, classifier = self.primary_classifier, model_type = "primary", ) # If primary result meets confidence threshold, return it if (primary_result.evidence_strength >= min_confidence): return primary_result # If primary is low confidence but we have fallback, try fallback if self.fallback_classifier: logger.info("Primary classifier low confidence, trying fallback model...") fallback_result = self._classify_with_model(text = text, classifier = self.fallback_classifier, model_type = "fallback", ) # Use fallback if it has higher confidence if (fallback_result.evidence_strength > primary_result.evidence_strength): return fallback_result # Return primary result even if low confidence return primary_result except Exception as e: logger.error(f"Error in domain classification: {repr(e)}") # Try fallback classifier if primary failed if self.fallback_classifier: try: logger.info("Trying fallback classifier after primary failure...") return self._classify_with_model(text = text, classifier = self.fallback_classifier, model_type = "fallback", ) except Exception as fallback_error: logger.error(f"Fallback classifier also failed: {repr(fallback_error)}") # Both models failed, return default return self._get_default_prediction() def _classify_with_model(self, text: str, classifier, model_type: str) -> DomainPrediction: """ Classify using a zero-shot classification model """ # Preprocess text processed_text = self._preprocess_text(text) # Get all candidate labels all_labels = list() label_to_domain = dict() for domain, labels in self.DOMAIN_LABELS.items(): # Use the first label as the primary one for this domain primary_label = labels[0] all_labels.append(primary_label) label_to_domain[primary_label] = domain # Perform zero-shot classification result = classifier(processed_text, candidate_labels = all_labels, multi_label = False, hypothesis_template = "This text is about {}.", ) # Convert to domain scores domain_scores = dict() for label, score in zip(result['labels'], result['scores']): domain = label_to_domain[label] domain_key = domain.value if (domain_key not in domain_scores): domain_scores[domain_key] = list() domain_scores[domain_key].append(score) # Average scores for each domain avg_domain_scores = {domain: sum(scores) / len(scores) for domain, scores in domain_scores.items()} # Sort by score sorted_domains = sorted(avg_domain_scores.items(), key = lambda x: x[1], reverse = True) # Get primary and secondary domains primary_domain_str, primary_score = sorted_domains[0] primary_domain = Domain(primary_domain_str) secondary_domain = None secondary_score = 0.0 # Use constant for secondary domain minimum score secondary_min_score = domain_classification_params.SECONDARY_DOMAIN_MIN_SCORE if ((len(sorted_domains) > 1) and (sorted_domains[1][1] >= secondary_min_score)): secondary_domain = Domain(sorted_domains[1][0]) secondary_score = sorted_domains[1][1] # Calculate evidence_strength evidence_strength = primary_score # Use constants for mixed domain detection high_conf_threshold = domain_classification_params.HIGH_CONFIDENCE_THRESHOLD mixed_secondary_min = domain_classification_params.MIXED_DOMAIN_SECONDARY_MIN mixed_ratio_thresh = domain_classification_params.MIXED_DOMAIN_RATIO_THRESHOLD mixed_conf_penalty = domain_classification_params.MIXED_DOMAIN_CONFIDENCE_PENALTY # If we have mixed domains with close scores, adjust confidence if (secondary_domain and (primary_score < high_conf_threshold) and (secondary_score > mixed_secondary_min)): score_ratio = secondary_score / primary_score # Secondary is at least 60% of primary if (score_ratio > mixed_ratio_thresh): # Lower confidence for mixed domains evidence_strength = ((primary_score + secondary_score) / 2 * mixed_conf_penalty) logger.info(f"Mixed domain detected: {primary_domain.value} + {secondary_domain.value}, will use interpolated thresholds") # Use constant for low confidence threshold low_conf_threshold = domain_classification_params.LOW_CONFIDENCE_THRESHOLD # If primary score is low and we have a secondary, it's uncertain if ((primary_score < low_conf_threshold) and secondary_domain): # Reduce confidence using penalty evidence_strength *= mixed_conf_penalty logger.info(f"{model_type.capitalize()} model classified domain: {primary_domain.value} (confidence: {evidence_strength:.3f})") return DomainPrediction(primary_domain = primary_domain, secondary_domain = secondary_domain, evidence_strength = evidence_strength, domain_scores = avg_domain_scores, ) def _preprocess_text(self, text: str) -> str: """ Preprocess text for classification """ # Truncate to reasonable length using constant max_words = domain_classification_params.MAX_WORDS_FOR_CLASSIFICATION words = text.split() if (len(words) > max_words): text = ' '.join(words[:max_words]) # Clean up text text = text.strip() if not text: return "general content" return text def _get_default_prediction(self) -> DomainPrediction: """ Get default prediction when classification fails """ return DomainPrediction(primary_domain = Domain.GENERAL, secondary_domain = None, evidence_strength = 0.5, domain_scores = {Domain.GENERAL.value: 1.0}, ) def get_adaptive_thresholds(self, domain_prediction: DomainPrediction): """ Get adaptive thresholds based on domain prediction """ # Use constants for threshold decisions high_conf_threshold = domain_classification_params.HIGH_CONFIDENCE_THRESHOLD med_conf_threshold = domain_classification_params.MEDIUM_CONFIDENCE_THRESHOLD # High confidence, single domain - use domain-specific thresholds if ((domain_prediction.evidence_strength > high_conf_threshold) and (not domain_prediction.secondary_domain)): return get_threshold_for_domain(domain_prediction.primary_domain) # Mixed domains - interpolate between primary and secondary if domain_prediction.secondary_domain: primary_score = domain_prediction.domain_scores.get(domain_prediction.primary_domain.value, 0) secondary_score = domain_prediction.domain_scores.get(domain_prediction.secondary_domain.value, 0) if (primary_score + secondary_score > 0): weight1 = primary_score / (primary_score + secondary_score) else: weight1 = domain_prediction.evidence_strength return interpolate_thresholds(domain1 = domain_prediction.primary_domain, domain2 = domain_prediction.secondary_domain, weight1 = weight1, ) # Low/medium confidence - blend with general domain if (domain_prediction.evidence_strength < med_conf_threshold): return interpolate_thresholds(domain1 = domain_prediction.primary_domain, domain2 = Domain.GENERAL, weight1 = domain_prediction.evidence_strength, ) # Default: use domain-specific thresholds return get_threshold_for_domain(domain_prediction.primary_domain) def cleanup(self): """ Clean up resources """ self.primary_classifier = None self.fallback_classifier = None self.is_initialized = False def quick_classify(text: str, **kwargs) -> DomainPrediction: """ Quick domain classification with default settings Arguments: ---------- text { str } : Input text **kwargs : Override settings Returns: -------- { DomainPrediction } : DomainPrediction object """ classifier = DomainClassifier() classifier.initialize() return classifier.classify(text, **kwargs) def get_domain_name(domain: Domain) -> str: """ Get human-readable domain name Arguments: ---------- domain { Domain } : Domain enum value Returns: -------- { str } : Human-readable domain name """ domain_names = {Domain.ACADEMIC : "Academic", Domain.CREATIVE : "Creative Writing", Domain.AI_ML : "AI/ML", Domain.SOFTWARE_DEV : "Software Development", Domain.TECHNICAL_DOC : "Technical Documentation", Domain.ENGINEERING : "Engineering", Domain.SCIENCE : "Science", Domain.BUSINESS : "Business", Domain.JOURNALISM : "Journalism", Domain.SOCIAL_MEDIA : "Social Media", Domain.BLOG_PERSONAL : "Personal Blog", Domain.LEGAL : "Legal", Domain.MEDICAL : "Medical", Domain.MARKETING : "Marketing", Domain.TUTORIAL : "Tutorial", Domain.GENERAL : "General", } return domain_names.get(domain, "Unknown") def is_technical_domain(domain: Domain) -> bool: """ Check if domain is technical in nature Arguments: ---------- domain { Domain } : Domain enum value Returns: -------- { bool } : True if technical domain """ technical_domains = {Domain.AI_ML, Domain.SOFTWARE_DEV, Domain.TECHNICAL_DOC, Domain.ENGINEERING, Domain.SCIENCE, } return domain in technical_domains def is_creative_domain(domain: Domain) -> bool: """ Check if domain is creative in nature Arguments: ---------- domain { Domain } : Domain enum value Returns: -------- { bool } : True if creative domain """ creative_domains = {Domain.CREATIVE, Domain.JOURNALISM, Domain.SOCIAL_MEDIA, Domain.BLOG_PERSONAL, Domain.MARKETING, } return domain in creative_domains def is_formal_domain(domain: Domain) -> bool: """ Check if domain is formal in nature Arguments: ---------- domain { Domain } : Domain enum value Returns: -------- { bool } : True if formal domain """ formal_domains = {Domain.ACADEMIC, Domain.LEGAL, Domain.MEDICAL, Domain.BUSINESS, } return domain in formal_domains # Export __all__ = ["Domain", "DomainClassifier", "DomainPrediction", "quick_classify", "get_domain_name", "is_technical_domain", "is_creative_domain", "is_formal_domain", ]