File size: 1,570 Bytes
b99688d | 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 | """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)
|