Spaces:
Running
Running
| """ | |
| Pré-filtre visuel CLIP (ViT-B/32) — triage zero-shot des documents KYC. | |
| Place dans le pipeline : APRES le televersement, AVANT Docling + Gemini. | |
| Objectif : economiser le quota Gemini en ecartant rapidement (CPU, < 1s) | |
| les documents qui ne ressemblent clairement pas au type attendu. | |
| Design retenu (cf. discussion architecture) : | |
| - Zero-shot : aucun dataset a constituer, aucun fine-tuning. Les classes | |
| sont decrites par des phrases (prompts) comparees aux embeddings d'image. | |
| - Triage, PAS rejet binaire : seuls les desaccords a haute confiance | |
| (>= SEUIL_BLOCAGE) bloquent. Une basse confiance repart vers le | |
| pipeline classique (Docling + Gemini) — un document atypique mais | |
| legitime ne doit jamais etre perdu. | |
| - INT8 dynamique sur CPU : quantification des couches Linear (~4x moins | |
| de memoire, accelération sur CPU). Fallback FP32 si indisponible. | |
| - PDF : rasterisation de la 1re page via PyMuPDF (aucune dependance | |
| systeme Poppler — compatible HF Spaces). | |
| Dependances a ajouter dans requirements.txt : | |
| open_clip_torch | |
| pymupdf | |
| torch (deja present) | |
| Pillow (deja present) | |
| """ | |
| import os | |
| import io | |
| import torch | |
| import open_clip | |
| from PIL import Image | |
| # Seuil de confiance a partir duquel un desaccord type_attendu/detecte | |
| # bloque le document AVANT de consommer du quota Gemini. | |
| # 0.80 = tolerant (recommande pour demarrer, a ajuster apres mesure). | |
| SEUIL_BLOCAGE = 0.80 | |
| # Libelles d'affichage (cles = identifiants internes = doc_type du formulaire) | |
| CATEGORY_LABELS = { | |
| "CNI_PASSEPORT": "CNI / Passeport", | |
| "JUSTIFICATIF_ADRESSE": "Justificatif d'adresse", | |
| "BULLETIN_SALAIRE": "Bulletin de paie", | |
| "RIB": "RIB / Releve bancaire", | |
| "RELEVE_COMPTE": "Relevé bancaire / Mobile Money (Wave, OM)", | |
| "AUTRE": "Autre document", | |
| } | |
| # Prompts textuels zero-shot : plusieurs par categorie pour robustesse. | |
| PROMPTS = { | |
| "CNI_PASSEPORT": [ | |
| "a photo of an identity card or a passport page with a portrait photo and official stamps", | |
| "a national identity card document", | |
| ], | |
| "JUSTIFICATIF_ADRESSE": [ | |
| "a utility bill, rent receipt or official proof of address showing a postal address", | |
| "an invoice or official letter showing a home address", | |
| ], | |
| "BULLETIN_SALAIRE": [ | |
| "a payslip or salary statement document with tables of amounts and deductions", | |
| "a pay stub with a net salary", | |
| ], | |
| "RIB": [ | |
| "a bank account details document with an IBAN number and a bank logo", | |
| "a bank statement header with bank coordinates", | |
| ], | |
| "RELEVE_COMPTE": [ | |
| "a bank statement or mobile money statement showing a list of transactions and account balance", | |
| "a mobile money account statement from Wave or Orange Money with transaction history", | |
| ], | |
| "AUTRE": [ | |
| "a random photo, a screenshot or any unrelated document", | |
| ], | |
| } | |
| class DocumentClassifier: | |
| """Singleton — le chargement du modele ne se fait qu'une fois.""" | |
| _instance = None | |
| def __new__(cls): | |
| if cls._instance is None: | |
| cls._instance = super(DocumentClassifier, cls).__new__(cls) | |
| cls._instance._load() | |
| return cls._instance | |
| def _load(self): | |
| print("[PREFILTRE] Chargement de CLIP ViT-B/32...") | |
| model, _, self.preprocess = open_clip.create_model_and_transforms( | |
| "ViT-B-32", pretrained="laion2b_s34b_b79k" | |
| ) | |
| # Quantification INT8 dynamique (CPU) — les deux tours separement, | |
| # inplace. Si l'un echoue, on continue en FP32 (budget < 1s tenu). | |
| try: | |
| torch.quantization.quantize_dynamic(model.visual, {torch.nn.Linear}, dtype=torch.qint8, inplace=True) | |
| except Exception as e: | |
| print(f"[PREFILTRE] Quantification 'visual' impossible ({e}) — FP32.") | |
| try: | |
| torch.quantization.quantize_dynamic(model.text, {torch.nn.Linear}, dtype=torch.qint8, inplace=True) | |
| except Exception as e: | |
| print(f"[PREFILTRE] Quantification 'text' impossible ({e}) — FP32.") | |
| model.eval() | |
| self.model = model | |
| self.tokenizer = open_clip.get_tokenizer("ViT-B-32") | |
| # Encodage des prompts : fait UNE SEULE FOIS au chargement. | |
| self.prompt_index = [] # position -> categorie | |
| all_prompts = [] | |
| for cat, prompts in PROMPTS.items(): | |
| for p in prompts: | |
| all_prompts.append(p) | |
| self.prompt_index.append(cat) | |
| with torch.no_grad(): | |
| tokens = self.tokenizer(all_prompts) | |
| txt = self.model.encode_text(tokens) | |
| self.text_features = txt / txt.norm(dim=-1, keepdim=True) | |
| print("[PREFILTRE] Pret.") | |
| def _to_image(file_bytes, file_name): | |
| """PDF -> rasterise la 1re page (PyMuPDF, sans Poppler). Sinon PIL.""" | |
| ext = os.path.splitext(file_name)[1].lower() | |
| if ext == ".pdf": | |
| import fitz # pymupdf | |
| doc = fitz.open(stream=file_bytes, filetype="pdf") | |
| pix = doc[0].get_pixmap(dpi=144) | |
| img = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB") | |
| doc.close() | |
| return img | |
| return Image.open(io.BytesIO(file_bytes)).convert("RGB") | |
| def classify(self, file_bytes, file_name): | |
| """ | |
| Retourne (categorie_detectee, confiance) avec categorie dans | |
| CATEGORY_LABELS. Confiance = probabilite softmax agregee par | |
| categorie (moyenne sur ses prompts). | |
| """ | |
| img = self.preprocess(self._to_image(file_bytes, file_name)).unsqueeze(0) | |
| feats = self.model.encode_image(img) | |
| feats = feats / feats.norm(dim=-1, keepdim=True) | |
| logits = self.model.logit_scale.exp() * feats @ self.text_features.T | |
| probs = logits.softmax(dim=-1).cpu().numpy()[0] | |
| agg = {} | |
| for p, cat in zip(probs, self.prompt_index): | |
| agg[cat] = agg.get(cat, 0.0) + float(p) | |
| total = sum(agg.values()) or 1.0 | |
| best = max(agg, key=agg.get) | |
| return best, agg[best] / total | |
| # ── Test rapide en ligne de commande ───────────────────────────────────────── | |
| if __name__ == "__main__": | |
| import sys | |
| clf = DocumentClassifier() | |
| for path in sys.argv[1:]: | |
| with open(path, "rb") as fh: | |
| cat, conf = clf.classify(fh.read(), os.path.basename(path)) | |
| print(f"{path} -> {CATEGORY_LABELS.get(cat, cat)} [{conf:.1%}]") |