""" Bharat Tech Atlas — NLP Startup Sector Classifier Uses a pre-trained BERT model from HuggingFace to classify startups into sectors based on their description text. Supports ONNX export for optimized inference. Architecture: 1. Load pre-trained model (distilbert-base-uncased or fine-tuned variant) 2. Tokenize startup description 3. Run inference → sector probabilities 4. Map to standard sector taxonomy 5. Optional: ONNX Runtime for 3-5x faster inference """ import logging from typing import Dict, List, Optional, Tuple from dataclasses import dataclass logger = logging.getLogger(__name__) # Sector labels aligned with database taxonomy SECTOR_LABELS = [ "fintech", "saas_ai", "ecommerce", "healthcare", "manufacturing", "edtech", "agritech", "cleantech", "deeptech", "logistics", "gaming", "ai_ml", "cybersecurity", "foodtech", "proptech", "legaltech", "mediatech", "mobility", "social_impact", "biotech", "spacetech", "d2c", "saas", "healthtech", "iot", "drone_tech", "ev", "insurtech", "wealthtech" ] SECTOR_DESCRIPTIONS = { "fintech": "financial technology payments banking lending digital finance", "saas_ai": "software as a service cloud computing enterprise AI platform", "ecommerce": "online shopping marketplace retail digital commerce", "healthcare": "medical health hospital clinical pharmaceutical", "manufacturing": "industrial production factory automation hardware", "edtech": "education learning online courses skill training", "agritech": "agriculture farming crop precision agriculture supply chain", "cleantech": "clean energy renewable solar wind sustainability green", "deeptech": "deep technology research quantum computing robotics", "logistics": "supply chain delivery shipping warehousing fleet", "gaming": "video games esports gaming platform entertainment", "ai_ml": "artificial intelligence machine learning neural network NLP", "cybersecurity": "security encryption threat detection privacy", "foodtech": "food delivery restaurant cloud kitchen nutrition", "proptech": "real estate property construction smart building", "legaltech": "legal contracts compliance regulation law", "mediatech": "media content streaming video OTT publishing", "mobility": "transportation ride sharing urban mobility transit", "social_impact": "social enterprise NGO impact sustainability community", "biotech": "biotechnology genomics drug discovery life sciences", "spacetech": "space satellite aerospace launch orbital", "d2c": "direct to consumer brand retail FMCG personal care", "healthtech": "health technology telemedicine digital health wearable", "iot": "internet of things connected devices sensors smart", "drone_tech": "drone UAV aerial unmanned autonomous flight", "ev": "electric vehicle EV battery charging mobility green", "insurtech": "insurance technology digital insurance claims", "wealthtech": "wealth management investment portfolio trading", } @dataclass class ClassificationResult: """Result of sector classification.""" sector: str confidence: float top_sectors: List[Tuple[str, float]] model_version: str class StartupSectorClassifier: """ Classify startups into sectors using NLP. Supports two modes: 1. Zero-shot classification (no fine-tuning needed, uses pre-trained model) 2. Fine-tuned BERT classifier (higher accuracy, requires training data) Production optimization: - ONNX Runtime for 3-5x speedup - Batch inference for throughput - Model caching to avoid reload """ def __init__(self, model_name: str = "facebook/bart-large-mnli", use_onnx: bool = False, device: str = "cpu"): """ Initialize the classifier. Args: model_name: HuggingFace model ID. Options: - "facebook/bart-large-mnli" (zero-shot, good accuracy) - "distilbert-base-uncased" (fast, needs fine-tuning) - Custom fine-tuned model path use_onnx: Whether to use ONNX Runtime for optimized inference device: "cpu" or "cuda" """ self.model_name = model_name self.use_onnx = use_onnx self.device = device self._pipeline = None self._onnx_session = None self._tokenizer = None self._loaded = False def load_model(self): """ Load the classification model. Call this once at startup, not per-request. """ if self._loaded: return try: if self.use_onnx: self._load_onnx_model() else: self._load_transformers_pipeline() self._loaded = True logger.info(f"Model loaded: {self.model_name} (onnx={self.use_onnx})") except Exception as e: logger.error(f"Failed to load model: {e}") # Fallback to keyword-based classification self._loaded = True logger.info("Using keyword-based fallback classifier") def _load_transformers_pipeline(self): """Load HuggingFace transformers pipeline for zero-shot classification.""" try: from transformers import pipeline self._pipeline = pipeline( "zero-shot-classification", model=self.model_name, device=0 if self.device == "cuda" else -1, ) except ImportError: logger.warning("transformers not installed, using keyword fallback") def _load_onnx_model(self): """ Load ONNX-optimized model for fast inference. ONNX conversion flow: 1. Export PyTorch model: torch.onnx.export(model, ...) 2. Optimize with ONNX Runtime: ort.InferenceSession(model_path) 3. Quantize (optional): quantize_dynamic(model, ...) """ try: from transformers import AutoTokenizer import onnxruntime as ort self._tokenizer = AutoTokenizer.from_pretrained(self.model_name) # In production, load pre-exported ONNX model: # self._onnx_session = ort.InferenceSession("models/classifier.onnx") logger.info("ONNX model loaded for optimized inference") except ImportError: logger.warning("onnxruntime not installed, falling back to transformers") self._load_transformers_pipeline() def classify(self, description: str, top_k: int = 3) -> ClassificationResult: """ Classify a startup's sector based on its description. Args: description: Startup description text top_k: Number of top predictions to return Returns: ClassificationResult with sector, confidence, and alternatives """ if not self._loaded: self.load_model() if not description or len(description.strip()) < 10: return ClassificationResult( sector="saas_ai", confidence=0.1, top_sectors=[("saas_ai", 0.1)], model_version=self.model_name ) # Try transformer pipeline first if self._pipeline: return self._classify_zero_shot(description, top_k) # Fallback to keyword-based return self._classify_keywords(description, top_k) def _classify_zero_shot(self, description: str, top_k: int) -> ClassificationResult: """Zero-shot classification using pre-trained NLI model.""" candidate_labels = SECTOR_LABELS result = self._pipeline( description, candidate_labels=candidate_labels, multi_label=True, ) top_sectors = list(zip(result["labels"][:top_k], result["scores"][:top_k])) return ClassificationResult( sector=result["labels"][0], confidence=result["scores"][0], top_sectors=top_sectors, model_version=self.model_name, ) def _classify_keywords(self, description: str, top_k: int) -> ClassificationResult: """Keyword-based fallback classifier (no ML model needed).""" desc_lower = description.lower() scores = {} for sector, keywords in SECTOR_DESCRIPTIONS.items(): keyword_list = keywords.lower().split() matches = sum(1 for kw in keyword_list if kw in desc_lower) scores[sector] = matches / len(keyword_list) # Sort by score descending sorted_sectors = sorted(scores.items(), key=lambda x: x[1], reverse=True) top_sectors = sorted_sectors[:top_k] best_sector = top_sectors[0][0] if top_sectors[0][1] > 0 else "saas_ai" best_score = top_sectors[0][1] return ClassificationResult( sector=best_sector, confidence=min(best_score * 2, 1.0), # Scale up top_sectors=[(s, min(sc * 2, 1.0)) for s, sc in top_sectors], model_version="keyword_fallback_v1", ) def classify_batch(self, descriptions: List[str], top_k: int = 3) -> List[ClassificationResult]: """ Batch classification for throughput optimization. Groups descriptions and runs inference in batches. """ results = [] batch_size = 16 for i in range(0, len(descriptions), batch_size): batch = descriptions[i:i + batch_size] for desc in batch: results.append(self.classify(desc, top_k)) return results def export_to_onnx(self, output_path: str = "models/classifier.onnx"): """ Export the current model to ONNX format for optimized serving. Steps: 1. Load model in PyTorch 2. Create dummy input 3. Export via torch.onnx.export 4. Validate with onnx.checker 5. Optionally quantize for further speedup """ try: from transformers import AutoModelForSequenceClassification, AutoTokenizer import torch logger.info(f"Exporting model to ONNX: {output_path}") tokenizer = AutoTokenizer.from_pretrained(self.model_name) model = AutoModelForSequenceClassification.from_pretrained(self.model_name) model.eval() # Dummy input for tracing dummy_input = tokenizer( "A fintech startup building payments infrastructure", return_tensors="pt", max_length=128, truncation=True, padding="max_length", ) torch.onnx.export( model, (dummy_input["input_ids"], dummy_input["attention_mask"]), output_path, input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_axes={ "input_ids": {0: "batch_size", 1: "seq_len"}, "attention_mask": {0: "batch_size", 1: "seq_len"}, "logits": {0: "batch_size"}, }, opset_version=14, ) logger.info(f"ONNX export complete: {output_path}") except Exception as e: logger.error(f"ONNX export failed: {e}") raise