Spaces:
Sleeping
Sleeping
zofiasmolenasana
Merge value+attribute labels, add header label, RAG eval improvements
d051588 unverified | """LayoutLMv3-based cell encoder for spreadsheet structure understanding. | |
| Wraps ``microsoft/layoutlmv3-base`` to produce per-cell embeddings that | |
| fuse text content with 2-D spatial layout information. Each spreadsheet | |
| cell is treated as a single "word" with its bounding box derived from | |
| row/col coordinates. | |
| Supports frozen (feature extraction) and fine-tunable modes. | |
| """ | |
| from __future__ import annotations | |
| from typing import Optional | |
| import torch | |
| import torch.nn as nn | |
| def _try_load_layoutlm(): | |
| """Lazy import to avoid hard dependency on transformers.""" | |
| try: | |
| from transformers import LayoutLMv3Model, LayoutLMv3Tokenizer | |
| return LayoutLMv3Model, LayoutLMv3Tokenizer | |
| except ImportError: | |
| return None, None | |
| class LayoutLMCellEncoder(nn.Module): | |
| """Encode spreadsheet cells using LayoutLMv3. | |
| Each cell's text and bounding box are passed through the model to | |
| produce a 768-d (base) or 1024-d (large) embedding per cell. | |
| Args: | |
| model_name: HuggingFace model identifier. | |
| frozen: If True, freeze all LayoutLM parameters (feature extraction). | |
| max_cells: Maximum cells per sheet (truncate if exceeded). | |
| """ | |
| MODEL_DIM = 768 # layoutlmv3-base | |
| def __init__( | |
| self, | |
| model_name: str = "microsoft/layoutlmv3-base", | |
| frozen: bool = True, | |
| max_cells: int = 512, | |
| ): | |
| super().__init__() | |
| LayoutLMv3Model, LayoutLMv3Tokenizer = _try_load_layoutlm() | |
| if LayoutLMv3Model is None: | |
| raise ImportError( | |
| "LayoutLMv3 requires `pip install transformers`. " | |
| "Install it to use the LayoutLM cell encoder." | |
| ) | |
| self.tokenizer = LayoutLMv3Tokenizer.from_pretrained(model_name) | |
| self.model = LayoutLMv3Model.from_pretrained(model_name) | |
| self.max_cells = max_cells | |
| if frozen: | |
| for param in self.model.parameters(): | |
| param.requires_grad = False | |
| self.model.eval() | |
| self._frozen = frozen | |
| def _cells_to_inputs( | |
| self, | |
| texts: list[str], | |
| rows: list[int], | |
| cols: list[int], | |
| max_row: int, | |
| max_col: int, | |
| ) -> dict[str, torch.Tensor]: | |
| """Convert cell data to LayoutLMv3 input format. | |
| Bounding boxes are normalized to [0, 1000] range (LayoutLM convention). | |
| Each cell occupies one "word" position. | |
| """ | |
| n = min(len(texts), self.max_cells) | |
| row_scale = 1000.0 / max(max_row, 1) | |
| col_scale = 1000.0 / max(max_col, 1) | |
| words = [] | |
| boxes = [] | |
| for i in range(n): | |
| word = texts[i][:50] if texts[i] else "[empty]" | |
| words.append(word) | |
| x0 = int(cols[i] * col_scale) | |
| y0 = int(rows[i] * row_scale) | |
| x1 = min(x0 + int(col_scale), 1000) | |
| y1 = min(y0 + int(row_scale), 1000) | |
| boxes.append([x0, y0, x1, y1]) | |
| encoding = self.tokenizer( | |
| words, | |
| boxes=boxes, | |
| return_tensors="pt", | |
| padding="max_length", | |
| truncation=True, | |
| max_length=self.max_cells + 2, | |
| is_split_into_words=True, | |
| ) | |
| return encoding | |
| def forward( | |
| self, | |
| texts: list[str], | |
| rows: list[int], | |
| cols: list[int], | |
| max_row: int = 100, | |
| max_col: int = 100, | |
| ) -> torch.Tensor: | |
| """Produce (N, 768) cell embeddings. | |
| Args: | |
| texts: Cell text values. | |
| rows: Cell row indices. | |
| cols: Cell column indices. | |
| max_row: Maximum row in the sheet. | |
| max_col: Maximum column in the sheet. | |
| Returns: | |
| Tensor of shape (N, 768) with per-cell embeddings. | |
| """ | |
| n = len(texts) | |
| inputs = self._cells_to_inputs(texts, rows, cols, max_row, max_col) | |
| device = next(self.model.parameters()).device | |
| inputs = {k: v.to(device) for k, v in inputs.items()} | |
| if self._frozen: | |
| with torch.no_grad(): | |
| outputs = self.model(**inputs) | |
| else: | |
| outputs = self.model(**inputs) | |
| hidden = outputs.last_hidden_state[0] | |
| word_ids = inputs.get("input_ids", None) | |
| cell_embeddings = [] | |
| token_idx = 1 | |
| for i in range(min(n, self.max_cells)): | |
| if token_idx < hidden.size(0) - 1: | |
| cell_embeddings.append(hidden[token_idx]) | |
| token_idx += 1 | |
| else: | |
| cell_embeddings.append(torch.zeros(self.MODEL_DIM, device=device)) | |
| while len(cell_embeddings) < n: | |
| cell_embeddings.append(torch.zeros(self.MODEL_DIM, device=device)) | |
| return torch.stack(cell_embeddings[:n]) | |
| def is_available() -> bool: | |
| """Check if LayoutLMv3 dependencies are installed.""" | |
| m, _ = _try_load_layoutlm() | |
| return m is not None | |