File size: 3,987 Bytes
81e8ada | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | from __future__ import annotations
import json
import os
from pathlib import Path
from typing import Any
import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoTokenizer
from encoder_model import EncoderClassifier
DEFAULT_REPO_ID = "ZenMan67/support-ticket-classifiers-minilm"
class HubTicketClassifier:
"""Download public checkpoints at startup and run the two-stage cascade."""
def __init__(
self,
repo_id: str | None = None,
revision: str = "main",
device: str | None = None,
preload_all: bool = False,
) -> None:
self.repo_id = repo_id or os.getenv("HF_MODEL_REPO", DEFAULT_REPO_ID)
self.revision = os.getenv("HF_MODEL_REVISION", revision)
self.device = torch.device(
device
or ("mps" if torch.backends.mps.is_available() else "cpu")
)
self.model_root = Path(snapshot_download(
repo_id=self.repo_id,
revision=self.revision,
repo_type="model",
allow_patterns=[
"tokenizer/*",
"handler/*",
],
))
self.tokenizer = AutoTokenizer.from_pretrained(
self.model_root / "tokenizer",
local_files_only=True,
)
self.models: dict[str, tuple[EncoderClassifier, dict[str, Any]]] = {}
self._load_task("handler")
if preload_all:
for task in ("human", "llm", "auto"):
self._load_task(task)
def _load_task(
self,
task: str,
) -> tuple[EncoderClassifier, dict[str, Any]]:
if task in self.models:
return self.models[task]
task_root = self.model_root / task
if not task_root.exists():
self.model_root = Path(snapshot_download(
repo_id=self.repo_id,
revision=self.revision,
repo_type="model",
allow_patterns=[f"{task}/*"],
))
task_root = self.model_root / task
config = json.loads(
(task_root / "config.json").read_text(encoding="utf-8")
)
model = EncoderClassifier(
config["base_model"],
config["num_labels"],
dropout=config["dropout"],
pooling=config["pooling"],
pretrained=False,
encoder_config=config["encoder_config"],
)
model.load_state_dict(load_file(task_root / "model.safetensors"))
model.to(self.device).eval()
self.models[task] = (model, config)
return model, config
@torch.inference_mode()
def classify(
self,
text: str,
task: str,
top_k: int = 3,
) -> dict[str, Any]:
model, config = self._load_task(task)
encoded = self.tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=config["max_length"],
)
logits = model(
encoded["input_ids"].to(self.device),
encoded["attention_mask"].to(self.device),
)
probabilities = logits.softmax(dim=-1)[0].cpu()
values, indices = probabilities.topk(
min(top_k, len(config["labels"]))
)
return {
"label": config["labels"][int(indices[0])],
"confidence": float(values[0]),
"top": [
{
"label": config["labels"][int(index)],
"confidence": float(value),
}
for value, index in zip(values, indices)
],
}
def predict(self, text: str, top_k: int = 3) -> dict[str, Any]:
handler = self.classify(text, "handler", top_k)
category = self.classify(text, handler["label"], top_k)
return {
"text": text,
"handler": handler,
"category": category,
}
|