| """Inferência do classificador: encoder multilíngue + regressão logística em JSON.""" |
| import json |
| import re |
|
|
| import numpy as np |
| from huggingface_hub import hf_hub_download |
| from sentence_transformers import SentenceTransformer |
|
|
|
|
| def chunks(text, size=900): |
| cur, out = "", [] |
| for p in text.split("\n"): |
| if len(cur) + len(p) > size and cur: |
| out.append(cur.strip()) |
| cur = "" |
| cur += p + "\n" |
| if cur.strip(): |
| out.append(cur.strip()) |
| return [c[:1400] for c in out] |
|
|
|
|
| class DocTypeClassifier: |
| def __init__(self, head, encoder=None): |
| self.classes = head["classes"] |
| self.coef = np.array(head["coef"]) |
| self.intercept = np.array(head["intercept"]) |
| self.mask = re.compile(head["label_mask"], re.I) |
| self.encoder = encoder or SentenceTransformer(head["encoder"], device="cpu") |
|
|
| def predict(self, text): |
| """Devolve [(classe, probabilidade)] da mais provável para a menos provável.""" |
| parts = chunks(self.mask.sub("[DOC]", text)) |
| v = self.encoder.encode(parts, normalize_embeddings=True, show_progress_bar=False).mean(axis=0) |
| v = v / np.linalg.norm(v) |
| scores = self.coef @ v + self.intercept |
| p = np.exp(scores - scores.max()) |
| p = p / p.sum() |
| return sorted(zip(self.classes, map(float, p)), key=lambda x: -x[1]) |
|
|
|
|
| def load(repo="fernandofa42/architecture-doc-type-classifier", encoder=None): |
| with open(hf_hub_download(repo, "head.json")) as fh: |
| return DocTypeClassifier(json.load(fh), encoder) |
|
|
