Automatic Document Classification by Type: ML Implementation

Implementation of Automatic Document Classification by Type

Document classification is the first stage in the Document Processing Pipeline. Before extracting data, the system must understand what document it's facing: an invoice, a waybill, a passport, or an act. Each type requires its own extractor, and an error at this stage breaks the entire chain. We implemented a multimodal classifier that simultaneously analyzes visual and textual features. This approach achieves accuracy of up to 97% on Russian documents — 5–7% higher than purely text-based solutions. Time savings on manual sorting reach 80%. Document processing cost reduction — up to 70%, with an average payback period of 6–12 months.

In this article, we'll break down the construction of a classification system, model selection, and typical mistakes. We'll describe a real implementation case. We guarantee quality at every stage — from labeling to deployment.

Why We Chose the Multimodal Approach?

Pure text classification based on OCR yields 85–90% accuracy. Adding visual features raises the bar to 94–97%. This is critical for documents with the same text but different layouts. For example, a bank invoice and a supplier invoice look different, though they contain similar fields. The visual analyzer captures logo placement, tables, and color blocks. The multimodal approach LayoutLMv3: Multi-modal Pre-training is 1.5 times more accurate on complex documents.

Multimodal Classification

The best approach is to use both visual and textual features simultaneously. Below is an example based on LayoutLMv3:

from transformers import LayoutLMv3ForSequenceClassification, LayoutLMv3Processor import torch import torch.nn as nn class DocumentClassifier: def __init__(self, model_path: str, doc_types: list[str]): self.processor = LayoutLMv3Processor.from_pretrained(model_path) self.model = LayoutLMv3ForSequenceClassification.from_pretrained( model_path, num_labels=len(doc_types) ) self.doc_types = doc_types self.model.eval() @torch.no_grad() def classify(self, image_path: str) -> dict: from PIL import Image image = Image.open(image_path).convert('RGB') encoding = self.processor( image, return_tensors='pt', truncation=True, max_length=512 ) outputs = self.model(**encoding) probs = torch.softmax(outputs.logits, dim=-1).squeeze() top_idx = probs.argmax().item() return { 'document_type': self.doc_types[top_idx], 'confidence': float(probs[top_idx]), 'all_scores': { self.doc_types[i]: float(probs[i]) for i in range(len(self.doc_types)) } } 

Training Without LayoutLM: EfficientNet + BERT

For a quick prototype without access to large models, we combine EfficientNet (visual encoder) and RuBERT (textual). This approach gives 91–94% accuracy on 15 classes — inferior to LayoutLMv3 but requires an order of magnitude fewer resources. The code is easily adaptable:

import timm from transformers import AutoTokenizer, AutoModel class LightweightDocClassifier(nn.Module): def __init__(self, num_classes: int): super().__init__() # Visual encoder self.visual = timm.create_model('efficientnet_b2', pretrained=True, num_classes=0) # Text encoder self.text_encoder = AutoModel.from_pretrained('DeepPavlov/rubert-base-cased') self.tokenizer = AutoTokenizer.from_pretrained('DeepPavlov/rubert-base-cased') # Fusion vis_dim = self.visual.num_features # 1408 text_dim = 768 self.fusion = nn.Sequential( nn.Linear(vis_dim + text_dim, 512), nn.GELU(), nn.Dropout(0.3), nn.Linear(512, num_classes) ) def forward(self, image_tensor, input_ids, attention_mask): vis_features = self.visual(image_tensor) text_out = self.text_encoder(input_ids, attention_mask) text_features = text_out.pooler_output # [CLS] token combined = torch.cat([vis_features, text_features], dim=-1) return self.fusion(combined) 

Typical Document Classes

Feature Collection for Classification

To improve accuracy, we add structural and textual features. They are especially useful for separating similar types, e.g., invoice vs. invoice (SF):

def extract_document_features(image_path: str, ocr_text: str) -> dict: return { # Structural features 'has_table': detect_tables(image_path), 'has_signature': detect_signature_zone(image_path), 'has_stamp': detect_stamp(image_path), 'has_photo': detect_person_photo(image_path), # Text patterns (regular expressions) 'has_inn': bool(re.search(r'\bИНН\b', ocr_text)), 'has_kpp': bool(re.search(r'\bКПП\b', ocr_text)), 'has_passport_series': bool(re.search(r'\d{4}\s\d{6}', ocr_text)), 'has_invoice_number': bool(re.search(r'№\s*\d+', ocr_text)), # Metadata 'aspect_ratio': get_aspect_ratio(image_path), 'orientation': detect_orientation(image_path), } 

How to Implement Classification in 5 Steps?

The implementation process is divided into clear stages:

1. Document flow analysis — study document types, volumes, sources, and current errors.
2. Dataset collection and labeling — collect at least 2000 samples per class, label the type. We use active learning to reduce costs.
3. Model selection and training — choose between LayoutLMv3 or EfficientNet+RuBERT based on analysis.
4. Validation and testing — test on real scans, measure Top-1 Accuracy and Macro F1.
5. Integration and deployment — package the model into a REST API, provide documentation.