sambodhan/urgency_classifier_space

1from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline2import torch3import os4 5class UrgencyPredictor:6    def __init__(self, model_repo="sambodhan/sambodhan_urgency_classifier",7                 cache_dir="/app/hf_cache"):8        """Load model and tokenizer once at startup."""9        10        self.model_repo = model_repo11        self.cache_dir = cache_dir12 13        # Ensure cache folder exists14        os.makedirs(self.cache_dir, exist_ok=True)15        16        # Device selection17        self.device = 0 if torch.cuda.is_available() else -118 19        print("Loading tokenizer and model...")20        # Load tokenizer and model21        self.tokenizer = AutoTokenizer.from_pretrained(self.model_repo, cache_dir=self.cache_dir, force_download=True)22        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_repo, cache_dir=self.cache_dir, force_download=True)23 24        # Create classification pipeline25        self.classifier = pipeline(26            "text-classification",27            model=self.model,28            tokenizer=self.tokenizer,29            device=self.device,30            return_all_scores=True31        )32        print("Model and tokenizer loaded successfully.")33 34    def predict(self, texts):35        """Predict urgency labels with scores for a single text or a batch."""36        if isinstance(texts, str):37            texts = [texts]38 39        results = self.classifier(texts)40        formatted_results = []41 42        for preds in results:43            # Sort by descending confidence44            preds = sorted(preds, key=lambda x: x["score"], reverse=True)45            top_pred = preds[0]46            label = top_pred["label"]47            confidence = round(top_pred["score"], 4)48            scores_dict = {p["label"]: round(p["score"], 4) for p in preds}49 50            formatted_results.append({51                "label": label,52                "confidence": confidence,53                "scores": scores_dict54            })55 56        # Return single dict if only one input57        return formatted_results[0] if len(formatted_results) == 1 else formatted_results58 59    @staticmethod60    def load_model():61        """Helper to preload the model during Docker build."""62        _ = UrgencyPredictor()63