Spaces:
Sleeping
Sleeping
| import re | |
| from typing import Dict, Any, List | |
| from ..data.schema import DocumentRecord, CLASS_FIELD_SCHEMAS | |
| INV_NUM_REGEX = re.compile(r'(?i)(?:invoice\s*(?:no|num|number|#)?[:.\s]*|inv\s*#?\s*|no[:.\s]+)([a-z0-9\-_/]+)') | |
| DATE_REGEX = re.compile(r'(\d{1,4}[/\-.]\d{1,2}[/\-.]\d{1,4})') | |
| AMOUNT_REGEX = re.compile(r'(?i)(?:total|amount|due|balance|subtotal)[\s:$]*([0-9]+[.,][0-9]{2,3})') | |
| class RuleBasedExtractor: | |
| """Baseline A: Rule-based regular expression and keyword heuristic extractor.""" | |
| def __init__(self): | |
| pass | |
| def extract(self, record: DocumentRecord, doc_class: str = None) -> Dict[str, Any]: | |
| cls_name = doc_class or record.doc_class | |
| tokens = [w.text for w in record.ocr_words] | |
| full_text = " ".join(tokens) | |
| schema_fields = CLASS_FIELD_SCHEMAS.get(cls_name, []) | |
| extracted = {f: None for f in schema_fields} | |
| if cls_name == "invoice": | |
| # 1. Invoice Number | |
| m_num = INV_NUM_REGEX.search(full_text) | |
| if m_num: | |
| extracted["invoice_number"] = m_num.group(1).strip() | |
| # 2. Invoice Date | |
| m_date = DATE_REGEX.search(full_text) | |
| if m_date: | |
| extracted["invoice_date"] = m_date.group(1).strip() | |
| # 3. Total Amount | |
| m_amt = AMOUNT_REGEX.search(full_text) | |
| if m_amt: | |
| extracted["total_amount"] = m_amt.group(1).strip() | |
| # 4. Buyer / Seller heuristics | |
| m_to = re.search(r'(?i)(?:bill\s*to|sold\s*to|to)[:\s]+([A-Za-z0-9\s]{3,30})', full_text) | |
| if m_to: | |
| extracted["buyer_name"] = m_to.group(1).strip() | |
| elif cls_name == "receipt": | |
| # 1. Company Name: top token block | |
| if tokens: | |
| extracted["company_name"] = " ".join(tokens[:min(3, len(tokens))]) | |
| # 2. Date | |
| m_date = DATE_REGEX.search(full_text) | |
| if m_date: | |
| extracted["date"] = m_date.group(1).strip() | |
| # 3. Total Amount | |
| m_amt = AMOUNT_REGEX.search(full_text) | |
| if m_amt: | |
| extracted["total_amount"] = m_amt.group(1).strip() | |
| elif cls_name == "application_form": | |
| # 1. Applicant Name | |
| m_name = re.search(r'(?i)(?:name|applicant|signature)[:\s]+([A-Za-z\s]{3,25})', full_text) | |
| if m_name: | |
| extracted["applicant_name"] = m_name.group(1).strip() | |
| # 2. Date | |
| m_date = DATE_REGEX.search(full_text) | |
| if m_date: | |
| extracted["date"] = m_date.group(1).strip() | |
| # 3. Address | |
| m_addr = re.search(r'(?i)(?:address|location|city)[:\s]+([A-Za-z0-9\s,]{5,40})', full_text) | |
| if m_addr: | |
| extracted["address"] = m_addr.group(1).strip() | |
| return extracted | |