from __future__ import annotations import hashlib import json import math import re from collections import Counter, defaultdict from dataclasses import dataclass from typing import Any _TOKEN_RE = re.compile(r"[A-Z0-9]{2,}") _STOP = { "THE","AND","FOR","FROM","WITH","THIS","THAT","OF","TO","IN","ON","AT","BY", "PVT","PRIVATE","LTD","LIMITED","LLP","INDIA","GST","GSTIN","INVOICE","BILL", "PAYMENT","PAID","RECEIPT","TRANSFER","BANK","ACCOUNT","AC","A","AN", } def _s(value) -> str: return str(value or "").strip() def amount_bucket(value) -> str: try: amount = abs(float(value or 0)) except Exception: amount = 0 if amount < 1_000: return "AMT_LT_1K" if amount < 10_000: return "AMT_1K_10K" if amount < 50_000: return "AMT_10K_50K" if amount < 200_000: return "AMT_50K_2L" if amount < 1_000_000: return "AMT_2L_10L" return "AMT_GE_10L" def text_tokens(value: str, prefix: str = "") -> list[str]: tokens = [] for token in _TOKEN_RE.findall(_s(value).upper()): if token in _STOP or len(token) < 2: continue if token.isdigit() and len(token) > 6: continue tokens.append(f"{prefix}{token}" if prefix else token) return tokens[:120] def feature_tokens(context: dict[str, Any]) -> list[str]: tokens: list[str] = [] source = _s(context.get("source_type")).upper() if source: tokens.append("SRC_" + source) direction = _s(context.get("direction")).upper() if direction: tokens.append("DIR_" + direction) hsn = re.sub(r"\D", "", _s(context.get("hsn_code"))) if hsn: tokens.append("HSN2_" + hsn[:2]) if len(hsn) >= 4: tokens.append("HSN4_" + hsn[:4]) tokens.extend(text_tokens(context.get("supplier_name", ""), "SUP_")) tokens.extend(text_tokens(context.get("supplier_gstin", ""), "GSTIN_")) tokens.extend(text_tokens(context.get("party_name", ""), "PTY_")) tokens.extend(text_tokens(context.get("description", ""), "TXT_")) tokens.extend(text_tokens(context.get("narration", ""), "TXT_")) tokens.extend(text_tokens(context.get("primary_industry", ""), "IND_")) tokens.extend(text_tokens(context.get("primary_business_activity", ""), "BUS_")) tokens.extend(text_tokens(context.get("main_products", ""), "PROD_")) tokens.extend(text_tokens(context.get("main_services", ""), "SERV_")) if context.get("inventory_maintained") is True: tokens.append("PROFILE_INVENTORY") if context.get("capital_intensive") is True: tokens.append("PROFILE_CAPITAL_INTENSIVE") if context.get("vehicle_intensive") is True: tokens.append("PROFILE_VEHICLE_INTENSIVE") if context.get("project_job_based") is True: tokens.append("PROFILE_PROJECT_BASED") tokens.append(amount_bucket(context.get("amount"))) # Preserve multiplicity modestly because NB benefits from repeated semantic cues. return tokens[:240] @dataclass class Prediction: class_id: int probability: float ranked: list[tuple[int, float]] evidence: list[str] class MultinomialNBClassifier: algorithm = "multinomial_nb_v1" @staticmethod def train(examples: list[dict[str, Any]], *, alpha: float = 1.0) -> dict[str, Any]: if len(examples) < 20: raise ValueError("At least 20 reviewed accounting examples are required to train the first internal model.") class_docs = Counter() class_token_counts: dict[int, Counter] = defaultdict(Counter) class_total_tokens = Counter() vocabulary = set() for example in examples: class_id = int(example["nature_id"]) features = feature_tokens(example["context"]) if not features: continue counts = Counter(features) class_docs[class_id] += 1 class_token_counts[class_id].update(counts) class_total_tokens[class_id] += sum(counts.values()) vocabulary.update(counts) classes = sorted(class_docs) if len(classes) < 2: raise ValueError("At least two different reviewed Accounting Natures are required for training.") usable_docs = sum(class_docs.values()) if usable_docs < 20: raise ValueError("At least 20 reviewed examples with usable accounting context are required.") model = { "algorithm": MultinomialNBClassifier.algorithm, "alpha": float(alpha), "document_count": usable_docs, "classes": classes, "class_docs": {str(k): int(v) for k, v in class_docs.items()}, "class_total_tokens": {str(k): int(v) for k, v in class_total_tokens.items()}, "token_counts": { str(class_id): dict(class_token_counts[class_id]) for class_id in classes }, "vocabulary_size": len(vocabulary), } return model @staticmethod def predict(model: dict[str, Any], context: dict[str, Any]) -> Prediction: features = Counter(feature_tokens(context)) if not features: raise ValueError("The transaction does not contain enough usable context for internal-model prediction.") classes = [int(value) for value in model["classes"]] class_docs = {int(k): int(v) for k, v in model["class_docs"].items()} class_total = {int(k): int(v) for k, v in model["class_total_tokens"].items()} token_counts = { int(k): {token: int(count) for token, count in values.items()} for k, values in model["token_counts"].items() } alpha = float(model.get("alpha") or 1.0) vocab_size = max(1, int(model.get("vocabulary_size") or 1)) total_docs = max(1, sum(class_docs.values())) class_count = len(classes) scores: dict[int, float] = {} evidence_by_class: dict[int, list[tuple[float, str]]] = defaultdict(list) for class_id in classes: prior = (class_docs[class_id] + alpha) / (total_docs + alpha * class_count) score = math.log(prior) denominator = class_total.get(class_id, 0) + alpha * vocab_size counts = token_counts.get(class_id, {}) for token, multiplicity in features.items(): numerator = counts.get(token, 0) + alpha contribution = multiplicity * math.log(numerator / denominator) score += contribution if counts.get(token, 0) > 0: evidence_by_class[class_id].append((counts[token], token)) scores[class_id] = score maximum = max(scores.values()) exp_scores = {key: math.exp(value - maximum) for key, value in scores.items()} total = sum(exp_scores.values()) or 1.0 probabilities = {key: value / total for key, value in exp_scores.items()} ranked = sorted(probabilities.items(), key=lambda item: (-item[1], item[0])) winner = ranked[0][0] evidence = [ token for _, token in sorted(evidence_by_class[winner], reverse=True)[:8] ] return Prediction( class_id=winner, probability=float(ranked[0][1]), ranked=[(int(k), float(v)) for k, v in ranked[:3]], evidence=evidence, ) def deterministic_split(examples: list[dict[str, Any]], validation_ratio: float = 0.20): train = [] validation = [] by_class: dict[int, list[dict[str, Any]]] = defaultdict(list) for example in examples: by_class[int(example["nature_id"])].append(example) for class_id, rows in by_class.items(): rows = sorted( rows, key=lambda row: hashlib.sha256( f"{row['source_type']}|{row['source_record_id']}|{class_id}".encode() ).hexdigest(), ) if len(rows) < 4: train.extend(rows) continue validation_count = max(1, int(round(len(rows) * validation_ratio))) validation.extend(rows[:validation_count]) train.extend(rows[validation_count:]) return train, validation