ZacharyZhang-NY
9cb46c18d3
The 500-line rule and the one-line-comment rule were enforced only from
memory, so they drifted. bench/filesize_check.py now checks the budget on
every tracked source file, and it passes.
Split, with behaviour proven unchanged by comparing each moved function body
against HEAD byte for byte (53 training bodies, 21 benchmark bodies, 0
differences):
db.py 503 -> 396 store_search.py 105
bench/run.py 609 -> 351 query_set.py 80, report.py 211
benchmarks.py 619 -> 424 benchmark_scoring.py 262
training.py 1630 -> 94 training_{common,encode,metrics,datasets,lora,retrieval}
TrainingService stays as a thin facade for the CLI; the facade's parameter
lists are pinned against the module functions by inspect, which is what
catches a leaked self.
Deleted the eight top-level functions no file called, after fixing the scan
that had reported 201 false positives by skipping the defining file. The
v1/v2/v3 sweep found nothing to delete and says so: every remaining version
tag is an on-disk format marker, a migration step, a fixture name, or a
third-party model name.
Collapsed 39 multi-line comment blocks (86 lines) to one line each, keeping
the measured facts and error strings that justify the code.
269 tests pass; ruff clean; mypy 24 errors in 9 files against a 24-in-6
baseline, net zero; pyright 4, the same four HEAD measures.
120 lines
4.9 KiB
Python
120 lines
4.9 KiB
Python
"""Image and text encoding for cross-modal retrieval, over one small ONNX graph.
|
|
|
|
Chinese-CLIP ViT-B/16 in fp16 (377 MB) keeps bilingual text and images in one vector
|
|
space, so a screenshot can be found by a picture or by a sentence. The graph is the
|
|
fused contrastive model: every call needs both towers' inputs, so the unused side is
|
|
fed a placeholder row.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
IMAGE_SIZE = 224
|
|
TEXT_MAX_LENGTH = 52
|
|
IMAGE_MEAN = (0.48145466, 0.45782750, 0.40821073)
|
|
IMAGE_STD = (0.26862954, 0.26130258, 0.27577711)
|
|
DOWNLOAD_HINT = (
|
|
"huggingface-cli download Xenova/chinese-clip-vit-base-patch16 "
|
|
"--include 'onnx/model_fp16.onnx' 'tokenizer*' 'config.json' --local-dir <dir>"
|
|
)
|
|
|
|
|
|
def preprocess_image(image: Image.Image) -> np.ndarray:
|
|
"""Turn a PIL image into the (3, 224, 224) float32 tensor the graph expects.
|
|
|
|
The reference processor resizes to a square and normalizes with CLIP constants;
|
|
doing it here keeps torchvision out of the dependency list.
|
|
"""
|
|
rgb = image.convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE), Image.Resampling.BICUBIC)
|
|
array = np.asarray(rgb, dtype=np.float32) / 255.0
|
|
mean = np.array(IMAGE_MEAN, dtype=np.float32)
|
|
std = np.array(IMAGE_STD, dtype=np.float32)
|
|
return ((array - mean) / std).transpose(2, 0, 1)
|
|
|
|
|
|
class ClipEncoder:
|
|
"""One ONNX session, both towers, batched image encoding."""
|
|
|
|
def __init__(self, model_dir: Path) -> None:
|
|
self.model_dir = Path(model_dir)
|
|
self.onnx_path = self.model_dir / "onnx" / "model_fp16.onnx"
|
|
if not self.onnx_path.exists():
|
|
raise SystemExit(
|
|
f"{self.onnx_path} is not a Chinese-CLIP export; fetch the fp16 graph: {DOWNLOAD_HINT}"
|
|
)
|
|
self._session: Any | None = None
|
|
self._tokenizer: Any | None = None
|
|
self._placeholder_text: dict[str, np.ndarray] | None = None
|
|
self._placeholder_image: np.ndarray | None = None
|
|
|
|
@property
|
|
def session(self) -> Any:
|
|
if self._session is None:
|
|
import onnxruntime as ort
|
|
|
|
self._session = ort.InferenceSession(str(self.onnx_path), providers=["CPUExecutionProvider"])
|
|
return self._session
|
|
|
|
@property
|
|
def tokenizer(self) -> Any:
|
|
if self._tokenizer is None:
|
|
from transformers import BertTokenizer
|
|
|
|
# The export is BERT-style (vocab.txt, [CLS]=101); transformers 5 aliases the fast and slow tokenizers, so this is the name the CLIP docs mean.
|
|
self._tokenizer = BertTokenizer.from_pretrained(str(self.model_dir))
|
|
return self._tokenizer
|
|
|
|
def encode_images(self, images: list[Image.Image], batch_size: int = 8) -> np.ndarray:
|
|
if not images:
|
|
return np.zeros((0, 512), dtype=np.float32)
|
|
rows: list[np.ndarray] = []
|
|
for start in range(0, len(images), batch_size):
|
|
chunk = images[start : start + batch_size]
|
|
pixels = np.stack([preprocess_image(image) for image in chunk])
|
|
rows.append(self._run("image_embeds", pixels, self._text_feed(1))[0])
|
|
return np.asarray(np.concatenate(rows), dtype=np.float32)
|
|
|
|
def encode_texts(self, texts: list[str], batch_size: int = 16) -> np.ndarray:
|
|
if not texts:
|
|
return np.zeros((0, 512), dtype=np.float32)
|
|
rows: list[np.ndarray] = []
|
|
for start in range(0, len(texts), batch_size):
|
|
chunk = texts[start : start + batch_size]
|
|
feed = self._text_feed(len(chunk), chunk)
|
|
rows.append(self._run("text_embeds", self._image_feed(len(chunk)), feed)[0])
|
|
return np.asarray(np.concatenate(rows), dtype=np.float32)
|
|
|
|
def _run(self, output: str, pixels: np.ndarray, text: dict[str, np.ndarray]) -> list[np.ndarray]:
|
|
feed = {"pixel_values": pixels.astype(np.float32), **text}
|
|
return self.session.run([output], feed)
|
|
|
|
def _text_feed(self, batch: int, texts: list[str] | None = None) -> dict[str, np.ndarray]:
|
|
if texts is None:
|
|
if self._placeholder_text is None:
|
|
self._placeholder_text = self._tokenize(["placeholder"])
|
|
return self._placeholder_text
|
|
return self._tokenize(texts)
|
|
|
|
def _image_feed(self, batch: int) -> np.ndarray:
|
|
if self._placeholder_image is None:
|
|
self._placeholder_image = np.zeros((1, 3, IMAGE_SIZE, IMAGE_SIZE), dtype=np.float32)
|
|
return np.repeat(self._placeholder_image, batch, axis=0)
|
|
|
|
def _tokenize(self, texts: list[str]) -> dict[str, np.ndarray]:
|
|
encoded = self.tokenizer(
|
|
texts,
|
|
padding="max_length",
|
|
truncation=True,
|
|
max_length=TEXT_MAX_LENGTH,
|
|
return_tensors="np",
|
|
)
|
|
return {
|
|
"input_ids": encoded["input_ids"].astype(np.int64),
|
|
"attention_mask": encoded["attention_mask"].astype(np.int64),
|
|
}
|