Spaces:
Sleeping
Sleeping
File size: 5,306 Bytes
c42f4ef | 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 | """Free, local models: text embedder, OpenCLIP image/text embedder, RapidOCR."""
from __future__ import annotations
import logging
from functools import cached_property
from threading import Lock
import numpy as np
from PIL import Image
from .config import Config
log = logging.getLogger(__name__)
def _l2norm(x: np.ndarray) -> np.ndarray:
x = np.asarray(x, dtype=np.float32)
n = np.linalg.norm(x, axis=-1, keepdims=True)
return x / np.clip(n, 1e-12, None)
class TextEmbedder:
"""Dense text embeddings. Backends: sentence-transformers (default BGE-small) or WordLlama."""
def __init__(self, cfg: Config):
self.cfg = cfg
self.backend = cfg.text_backend
if self.backend == "sentence-transformers":
from sentence_transformers import SentenceTransformer
self._m = SentenceTransformer(cfg.text_model, device=cfg.device)
get_dim = getattr(self._m, "get_embedding_dimension", None) or self._m.get_sentence_embedding_dimension
self.dim = int(get_dim())
self.name = f"st:{cfg.text_model}"
elif self.backend == "wordllama":
from wordllama import WordLlama
self._m = WordLlama.load()
self.dim = int(self._m.embed(["probe"]).shape[-1])
self.name = "wordllama:default"
else:
raise ValueError(f"unknown text_backend {self.backend!r}")
def _encode(self, texts: list[str]) -> np.ndarray:
if not texts:
return np.zeros((0, self.dim), np.float32)
if self.backend == "sentence-transformers":
v = self._m.encode(texts, batch_size=32, normalize_embeddings=True, show_progress_bar=False)
else:
v = self._m.embed(texts)
return _l2norm(v)
def embed_docs(self, texts: list[str]) -> np.ndarray:
return self._encode(texts)
def embed_query(self, text: str) -> np.ndarray:
prefix = self.cfg.query_prefix if self.backend == "sentence-transformers" else ""
return self._encode([prefix + text])[0]
class ClipEmbedder:
"""OpenCLIP: images and text in one 512-d space (text->image retrieval)."""
def __init__(self, cfg: Config):
import open_clip
import torch
self.torch = torch
self.device = cfg.device
self.model, _, self.preprocess = open_clip.create_model_and_transforms(
cfg.clip_model, pretrained=cfg.clip_pretrained, device=cfg.device
)
self.model.eval()
self.tokenizer = open_clip.get_tokenizer(cfg.clip_model)
self.dim = int(self.embed_text(["probe"]).shape[-1])
self.name = f"clip:{cfg.clip_model}:{cfg.clip_pretrained.split('/')[-1]}"
def embed_images(self, images: list[Image.Image], batch: int = 16) -> np.ndarray:
out = []
with self.torch.no_grad():
for i in range(0, len(images), batch):
x = self.torch.stack([self.preprocess(im.convert("RGB")) for im in images[i:i + batch]])
out.append(self.model.encode_image(x.to(self.device)).float().cpu().numpy())
return _l2norm(np.concatenate(out)) if out else np.zeros((0, self.dim), np.float32)
def embed_text(self, texts: list[str]) -> np.ndarray:
with self.torch.no_grad():
tok = self.tokenizer(texts).to(self.device)
return _l2norm(self.model.encode_text(tok).float().cpu().numpy())
class OCR:
"""RapidOCR (PaddleOCR models exported to ONNX; bundled in the wheel, no download)."""
def __init__(self, min_score: float = 0.5):
from rapidocr_onnxruntime import RapidOCR
self._engine = RapidOCR()
self.min_score = min_score
def __call__(self, image: Image.Image) -> str:
arr = np.array(image.convert("RGB"))[:, :, ::-1] # RGB -> BGR (OpenCV convention)
result, _ = self._engine(arr)
if not result:
return ""
# sort top-to-bottom, left-to-right
result.sort(key=lambda r: (round(r[0][0][1] / 12), r[0][0][0]))
return " ".join(t for _, t, s in result if float(s) >= self.min_score).strip()
class Models:
"""Lazy holder so nothing heavy loads until it is needed."""
def __init__(self, cfg: Config):
self.cfg = cfg
self._ocr_lock = Lock()
@cached_property
def text(self) -> TextEmbedder:
log.info("loading text embedder (%s)", self.cfg.text_backend)
return TextEmbedder(self.cfg)
@cached_property
def clip(self) -> ClipEmbedder:
log.info("loading OpenCLIP %s/%s", self.cfg.clip_model, self.cfg.clip_pretrained)
return ClipEmbedder(self.cfg)
@cached_property
def ocr(self) -> OCR | None:
return OCR() if self.cfg.ocr else None
def ocr_text(self, key: str, image: Image.Image) -> str:
"""OCR with a content-hash cache (same image in many sessions/files is read once)."""
if self.ocr is None:
return ""
# The OCR engine and its cache are shared by every browser session.
with self._ocr_lock:
cache = self.__dict__.setdefault("_ocr_cache", {})
if key not in cache:
if len(cache) >= 512:
cache.clear()
cache[key] = self.ocr(image)
return cache[key]
|