File size: 5,025 Bytes
82f262a | 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 | """
DocumentModule - PDF/DOCX/OCR with layout-aware parsing for MORPH-AI v6.
"""
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from architecture import MorphConfig
@dataclass
class DocumentFacts:
text: str = ""
pages: int = 0
format: str = ""
tables: List[List[List[str]]] = field(default_factory=list)
metadata: Dict[str, str] = field(default_factory=dict)
embedding: Optional[torch.Tensor] = None
def to_text(self) -> str:
parts = [f"document {self.format} {self.pages}p"]
if self.text:
parts.append(f"text: {self.text[:500]}")
if self.tables:
parts.append(f"tables: {len(self.tables)}")
return " | ".join(parts)
def to_dict(self) -> Dict[str, Any]:
return {
"text": self.text,
"pages": self.pages,
"format": self.format,
"tables": self.tables,
"metadata": self.metadata,
}
class DocumentModule(nn.Module):
"""PDF/DOCX/OCR with layout-aware parsing for document understanding."""
def __init__(self, config: MorphConfig, hidden_dim: int):
super().__init__()
self.max_pages = config.doc_max_pages
self.page_proj = nn.Linear(hidden_dim, config.doc_hidden)
self.layout_encoder = nn.Sequential(
nn.Linear(config.doc_hidden + 4, config.doc_hidden),
nn.GELU(),
nn.Linear(config.doc_hidden, hidden_dim),
)
nn.init.zeros_(self.layout_encoder[-1].weight)
nn.init.zeros_(self.layout_encoder[-1].bias)
def forward(self, hidden: torch.Tensor, layout_info: Optional[torch.Tensor] = None) -> torch.Tensor:
B, T, H = hidden.shape
page_emb = self.page_proj(hidden)
if layout_info is not None:
layout = layout_info.to(hidden.dtype)
page_emb = self.layout_encoder(torch.cat([page_emb, layout], dim=-1))
return hidden + page_emb
def extract_text(self, source) -> str:
"""Extract text from PDF/DOCX/image with OCR fallback."""
try:
if hasattr(source, 'endswith'):
if source.endswith('.pdf'):
return self._extract_pdf(source)
elif source.endswith('.docx'):
return self._extract_docx(source)
elif source.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
return self._extract_image_ocr(source)
return self._extract_image_ocr(source)
except Exception as e:
return f"[document extraction error: {e}]"
def _extract_pdf(self, path: str) -> str:
try:
import fitz
doc = fitz.open(path)
pages = []
for i in range(min(len(doc), self.max_pages)):
page = doc[i]
text = page.get_text()
tables = page.find_tables()
if tables.tables:
for table in tables.tables:
pages.append(f"[TABLE]\n{table.to_pandas().to_string()}")
pages.append(text)
return "\n\n".join(pages)
except ImportError:
return "[PDF extraction requires PyMuPDF: pip install pymupdf]"
def _extract_docx(self, path: str) -> str:
try:
import docx2txt
return docx2txt.process(path)
except ImportError:
return "[DOCX extraction requires docx2txt: pip install docx2txt]"
def _extract_image_ocr(self, source) -> str:
try:
import pytesseract
from PIL import Image
img = Image.open(source)
text = pytesseract.image_to_string(img)
data = pytesseract.image_to_data(img, output_type=pytesseract.Output.DICT)
lines = []
for i, word in enumerate(data['text']):
if word.strip():
lines.append(word)
return text + "\n\n[LAYOUT]\n" + " ".join(lines)
except ImportError:
return "[OCR requires pytesseract + Pillow: pip install pytesseract pillow]"
except Exception as e:
return f"[OCR error: {e}]"
def analyze(self, source) -> DocumentFacts:
"""Full document analysis returning structured facts."""
facts = DocumentFacts()
try:
text = self.extract_text(source)
facts.text = text
facts.pages = len(text.split('\n\n'))
if hasattr(source, 'endswith'):
if source.endswith('.pdf'):
facts.format = 'PDF'
elif source.endswith('.docx'):
facts.format = 'DOCX'
else:
facts.format = 'IMAGE'
else:
facts.format = 'UNKNOWN'
except Exception as e:
facts.text = f"[analysis error: {e}]"
return facts
|