Spaces:
Runtime error
Runtime error
File size: 7,655 Bytes
410242f | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | """
Document Classifier Agent - Classifies documents into categories.
Uses machine learning and heuristics for accurate classification.
"""
import logging
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from enum import Enum
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline
logger = logging.getLogger(__name__)
class DocumentType(str, Enum):
"""Supported document types"""
INVOICE = "invoice"
RECEIPT = "receipt"
CONTRACT = "contract"
REPORT = "report"
EMAIL = "email"
FORM = "form"
LETTER = "letter"
UNKNOWN = "unknown"
@dataclass
class ClassificationResult:
"""Result from document classification"""
document_type: DocumentType
confidence: float
probabilities: Dict[str, float]
relevant_keywords: List[str]
metadata: Dict
class DocumentClassifier:
"""Classify documents into predefined categories"""
# Keywords for different document types
KEYWORDS = {
DocumentType.INVOICE: [
'invoice', 'bill', 'amount due', 'invoice number',
'tax', 'subtotal', 'total amount', 'payment terms'
],
DocumentType.RECEIPT: [
'receipt', 'paid', 'purchase', 'amount paid',
'transaction', 'date', 'total', 'items'
],
DocumentType.CONTRACT: [
'agreement', 'contract', 'terms', 'conditions',
'hereby', 'whereas', 'party', 'signature'
],
DocumentType.REPORT: [
'report', 'summary', 'analysis', 'findings',
'conclusion', 'recommendation', 'data', 'chart'
],
DocumentType.EMAIL: [
'from:', 'to:', 'subject:', 'cc:', 'date:',
'regards', 'sincerely', 'dear', 'hello'
],
DocumentType.FORM: [
'form', 'field', 'checkbox', 'signature line',
'required', 'optional', 'please enter', 'x', '☑'
],
DocumentType.LETTER: [
'letter', 'dear', 'yours truly', 'sincerely',
'regards', 'yours', 'faithfully'
],
}
def __init__(self):
"""Initialize classifier with pre-trained models"""
self.vectorizer = TfidfVectorizer(max_features=100, lowercase=True)
self.classifier = None
self._initialize_classifier()
def _initialize_classifier(self):
"""Initialize ML classifier"""
try:
# Training data (simplified for demo)
training_docs = [
("Invoice #123 amount due $500", DocumentType.INVOICE),
("Receipt for purchase $45.99", DocumentType.RECEIPT),
("This agreement is entered into", DocumentType.CONTRACT),
("Monthly report and findings", DocumentType.REPORT),
]
texts = [doc[0] for doc in training_docs]
labels = [doc[1].value for doc in training_docs]
# Create pipeline
self.classifier = Pipeline([
('tfidf', TfidfVectorizer(max_features=100, lowercase=True)),
('clf', MultinomialNB())
])
self.classifier.fit(texts, labels)
logger.info("Document classifier initialized")
except Exception as e:
logger.warning(f"Classifier initialization failed: {e}")
def classify(self, text: str, metadata: Optional[Dict] = None) -> ClassificationResult:
"""
Classify a document based on its text content.
Args:
text: Document text to classify
metadata: Optional metadata (filename, file type, etc.)
Returns:
ClassificationResult with document type and confidence
"""
try:
if not text or not text.strip():
return ClassificationResult(
document_type=DocumentType.UNKNOWN,
confidence=0.0,
probabilities={},
relevant_keywords=[],
metadata=metadata or {}
)
# Keyword-based classification
keyword_scores = self._score_by_keywords(text)
best_type = max(keyword_scores, key=keyword_scores.get)
best_confidence = keyword_scores[best_type]
# ML-based classification (if available)
ml_type, ml_confidence = self._classify_with_ml(text)
# Combine results (weighted average)
if best_confidence > 0.3:
final_type = best_type
final_confidence = best_confidence * 0.7 + ml_confidence * 0.3
else:
final_type = ml_type
final_confidence = ml_confidence
# Extract relevant keywords
relevant_keywords = self._extract_relevant_keywords(text, final_type)
# Prepare probabilities
probabilities = {doc_type.value: score for doc_type, score in keyword_scores.items()}
return ClassificationResult(
document_type=final_type,
confidence=min(final_confidence, 1.0),
probabilities=probabilities,
relevant_keywords=relevant_keywords,
metadata=metadata or {}
)
except Exception as e:
logger.error(f"Classification failed: {str(e)}")
return ClassificationResult(
document_type=DocumentType.UNKNOWN,
confidence=0.0,
probabilities={},
relevant_keywords=[],
metadata=metadata or {}
)
def _score_by_keywords(self, text: str) -> Dict[DocumentType, float]:
"""Score document by keyword matching"""
text_lower = text.lower()
scores = {}
for doc_type, keywords in self.KEYWORDS.items():
matches = sum(1 for keyword in keywords if keyword in text_lower)
score = matches / len(keywords) if keywords else 0
scores[doc_type] = score
return scores
def _classify_with_ml(self, text: str) -> Tuple[DocumentType, float]:
"""Classify using ML model"""
try:
if self.classifier:
prediction = self.classifier.predict([text])[0]
probabilities = self.classifier.predict_proba([text])[0]
confidence = float(np.max(probabilities))
doc_type = DocumentType(prediction)
return doc_type, confidence
except Exception as e:
logger.warning(f"ML classification failed: {e}")
return DocumentType.UNKNOWN, 0.0
def _extract_relevant_keywords(self, text: str, doc_type: DocumentType) -> List[str]:
"""Extract keywords relevant to detected document type"""
text_lower = text.lower()
keywords = self.KEYWORDS.get(doc_type, [])
found_keywords = [kw for kw in keywords if kw in text_lower]
return found_keywords
def batch_classify(self, documents: List[Dict]) -> List[ClassificationResult]:
"""Classify multiple documents"""
results = []
for doc in documents:
text = doc.get('text', '')
metadata = doc.get('metadata')
result = self.classify(text, metadata)
results.append(result)
return results
|