seai2526-uniba-TheClouds/Code-Comment-Classification-Api

1"""Prediction helpers for different model types.2 3This module provides `ModelPredictor`, a lightweight wrapper that unifies4inference for SetFit, scikit-learn RandomForest pipelines, and HuggingFace5transformer sequence classification models. It standardizes inputs/outputs6to a NumPy array of shape (n_samples, n_labels).7"""8 9import os10from typing import List, Union11 12import joblib13import numpy as np14from setfit import SetFitModel15import torch16from transformers import AutoModelForSequenceClassification, AutoTokenizer17 18TextInput = Union[str, List[str]]19 20 21class ModelPredictor:22    """Unified predictor for SetFit, Random Forest and Transformer models.23 24    Expected directory layout:25 26        models/27          ├── java/28          │     ├── setfit/                # SetFit saved model directory29          │     ├── random_forest.joblib   # sklearn pipeline30          │     └── transformer/           # HF model + tokenizer (config.json, etc.)31          ├── python/32          │     ├── setfit/33          │     ├── random_forest.joblib34          │     └── transformer/35          └── pharo/36                ├── setfit/37                ├── random_forest.joblib38                └── transformer/39    """40 41    def __init__(42        self,43        lang: str,44        model_type: str,45        model_root: str = "models",46        threshold: float = 0.5,47        max_length: int = 128,48    ) -> None:49        """Parameters50 51        ----------52        lang : str53            One of {"java", "python", "pharo"}.54        model_type : str55            One of {"setfit", "random_forest", "transformer"}.56        model_root : str57            Root directory where models are stored.58        threshold : float59            Decision threshold for multi-label Transformer predictions.60            Ignored for SetFit and Random Forest (they already output labels).61        max_length : int62            Max sequence length for Transformer tokenization.63 64        """65        self.lang = lang66        self.model_type = model_type67        self.model_root = model_root68        self.threshold = float(threshold)69        self.max_length = int(max_length)70 71        # device only matters for Transformer72        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")73 74        if model_type == "setfit":75            model_path = os.path.join(self.model_root, self.lang, "setfit")76            if not os.path.isdir(model_path):77                raise FileNotFoundError(f"SetFit model not found at: {model_path}")78            self.model = SetFitModel.from_pretrained(model_path)79 80        elif model_type == "random_forest":81            model_path = os.path.join(self.model_root, self.lang, "random_forest.joblib")82            if not os.path.isfile(model_path):83                raise FileNotFoundError(f"Random Forest model not found at: {model_path}")84            self.model = joblib.load(model_path)85 86        elif model_type == "transformer":87            model_path = os.path.join(self.model_root, self.lang, "transformer")88            if not os.path.isdir(model_path):89                raise FileNotFoundError(f"Transformer model not found at: {model_path}")90 91            # load tokenizer and model from the same directory used during training92            self.tokenizer = AutoTokenizer.from_pretrained(model_path)93            self.model = AutoModelForSequenceClassification.from_pretrained(model_path).to(94                self.device95            )96            self.model.eval()97 98        else:99            raise ValueError(f"Unsupported model_type: {model_type}")100 101    def predict(self, texts: TextInput) -> np.ndarray:102        """Run prediction on one or many text samples.103 104        Parameters105        ----------106        texts : str | list[str]107            A single text or a list of texts.108 109        Returns110        -------111        np.ndarray112            Array of shape (n_samples, n_labels) with integer (typically binary) values.113 114        """115        if isinstance(texts, str):116            texts = [texts]117 118        if self.model_type == "setfit":119            raw_outputs = self.model(texts)120            outputs = np.array(list(raw_outputs), dtype=int)121 122        elif self.model_type == "random_forest":123            raw_outputs = self.model.predict(texts)124            outputs = np.array(list(raw_outputs), dtype=int)125 126        elif self.model_type == "transformer":127            enc = self.tokenizer(128                texts,129                padding=True,130                truncation=True,131                max_length=self.max_length,132                return_tensors="pt",133            )134            enc = {k: v.to(self.device) for k, v in enc.items()}135 136            with torch.no_grad():137                logits = self.model(**enc).logits138                probs = torch.sigmoid(logits)139                preds = (probs > self.threshold).long().cpu().numpy()140 141            outputs = preds.astype(int)142        else:143            raise ValueError(f"Unsupported model_type: {self.model_type}")144 145        # Ensure 2D shape (n_samples, n_labels)146        if outputs.ndim == 1:147            outputs = outputs.reshape(1, -1)148 149        return outputs150