fgonon/stackoverflow_tagpredict

1import torch2import json3import re4import os5import shutil6from pathlib import Path7from transformers import AutoTokenizer, AutoModelForSequenceClassification8from huggingface_hub import hf_hub_download9 10class TagPredictor:11    def __init__(self, model_path='model_stage2', tags_path='tags.json'):12        self.model_path = model_path13        self.tags_path = tags_path14        self.model = None15        self.tokenizer = None16        self.labels = None17        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')18        self._load_resources()19 20    def _load_resources(self):21        self._ensure_model_weights()22        # load tag names from json dictionary23        try:24            with open(self.tags_path, 'r') as f:25                # keys in json are always strings, convert to int for mapping if needed26                # but label_dict keys match these IDs27                raw_tags = json.load(f)28                tag_id_to_name = {int(k): v for k, v in raw_tags.items()}29        except FileNotFoundError:30            print(f"Warning: {self.tags_path} not found. Tags will be unknown.")31            tag_id_to_name = {}32 33        # load model and tokenizer34        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)35        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_path)36        self.model.to(self.device)37        self.model.eval()38 39        # load label dictionary (maps model class index -> tag_id)40        with open(f'{self.model_path}/label_dict.json', 'r') as f:41            label_dict = json.load(f)42            # map: class_index (value in label_dict) -> tag_name (via key in label_dict which is tag_id)43            self.labels = {v: tag_id_to_name.get(int(k), f'tag_{k}') for k, v in label_dict.items()}44 45    def _ensure_model_weights(self):46        weights_path = Path(self.model_path) / 'model.safetensors'47        if not weights_path.exists() or self._is_lfs_pointer(weights_path):48            self._download_weights(weights_path)49 50    def _is_lfs_pointer(self, file_path):51        try:52            with open(file_path, 'rb') as f:53                head = f.read(64)54            return head.startswith(b'version https://git-lfs.github.com/spec/v1')55        except OSError:56            return True57 58    def _download_weights(self, destination):59        repo_id = os.getenv('MODEL_REPO_ID') or os.getenv('SPACE_ID')60        if not repo_id:61            raise RuntimeError('model.safetensors is missing. set MODEL_REPO_ID or SPACE_ID so it can be downloaded.')62        repo_type = os.getenv('MODEL_REPO_TYPE', 'space')63        remote_path = f'{self.model_path}/model.safetensors'64        try:65            downloaded_path = hf_hub_download(repo_id=repo_id, filename=remote_path, repo_type=repo_type)66        except Exception as exc:67            raise RuntimeError(f'could not download {remote_path} from {repo_id}') from exc68        destination.parent.mkdir(parents=True, exist_ok=True)69        shutil.copy(downloaded_path, destination)70 71    def preprocess_text(self, text):72        text = text.lower()73        text = re.sub(r'http\S+', '', text)74        text = re.sub(r'[^a-z\s]', '', text)75        return text76 77    def predict(self, text, top_k=5):78        clean_text = self.preprocess_text(text)79        inputs = self.tokenizer(clean_text, truncation=True, padding='max_length', max_length=128, return_tensors='pt')80        inputs = {k: v.to(self.device) for k, v in inputs.items()}81 82        with torch.no_grad():83            outputs = self.model(**inputs)84            probs = torch.nn.functional.softmax(outputs.logits, dim=1)85            top_probs, top_ids = torch.topk(probs, top_k, dim=1)86 87        results = []88        for idx, prob in zip(top_ids[0], top_probs[0]):89            results.append({90                "tag": self.labels.get(idx.item(), 'unknown'),91                "probability": float(prob.item())92            })93        return results94 95# global instance for reuse96predictor = TagPredictor()97