Spaces:
No application file
No application file
File size: 6,283 Bytes
96d8696 | 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 145 146 147 148 149 150 151 152 153 154 155 156 | #!/usr/bin/env python
"""Stage 02 — extract embeddings.
The one GPU-heavy stage: runs the frozen EmbeddingGemma + SigLIP2 encoders
over every row and caches the results as .npy files, so training (stage 03)
never has to re-run the encoders.
Order of operations (deliberate, not incidental):
1. Load the text encoder, extract text embeddings for train AND test.
2. Free the text encoder from RAM/VRAM.
3. Load the image encoder, extract image embeddings for train AND test.
4. Free the image encoder.
This keeps peak memory to "one encoder at a time" instead of both loaded
simultaneously, and means a crash during image extraction doesn't force
re-running text extraction too.
Every batch is checkpointed to disk (see src/encoders/batch_extract.py) —
if this script crashes or the session times out, re-running the exact same
command resumes from the last checkpoint instead of starting over.
Usage:
python scripts/02_extract_embeddings.py --config configs/base.yaml
"""
import argparse
import sys
from pathlib import Path
import numpy as np
import pandas as pd
from PIL import Image
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.encoders.batch_extract import free_encoder, resumable_batch_encode
from src.encoders.registry import build_image_encoder, build_text_encoder
from src.utils.config import load_config
from src.utils.exceptions import PricePredictorError
from src.utils.logging import get_logger
from src.utils.seed import set_seed
logger = get_logger(__name__)
SPLITS = ("train", "test")
def _load_split_df(processed_dir: Path, split: str) -> pd.DataFrame:
parquet_path = processed_dir / f"{split}_clean.parquet"
if not parquet_path.exists():
raise PricePredictorError(f"{parquet_path} not found — run scripts/01_preprocess.py first")
return pd.read_parquet(parquet_path)
def run(config_path: str) -> None:
config = load_config(config_path)
set_seed(config["seed"])
data_cfg = config["data"]
processed_dir = Path(data_cfg["processed_dir"])
embeddings_dir = Path(data_cfg["embeddings_dir"])
embeddings_dir.mkdir(parents=True, exist_ok=True)
dfs = {split: _load_split_df(processed_dir, split) for split in SPLITS}
for split, df in dfs.items():
logger.info("Loaded split=%s: %d rows", split, len(df))
df[["sample_id"]].to_csv(embeddings_dir / f"{split}_ids.csv", index=False)
if split == "train":
np.save(embeddings_dir / "train_price.npy", df["price"].to_numpy(dtype=np.float32))
# ---- Phase 1: text embeddings (train + test), then free the encoder ----
logger.info("=== Phase 1/2: text embeddings ===")
text_cfg = config["encoders"]["text"]
text_encoder = build_text_encoder(text_cfg)
text_output_dim = text_cfg.get("mrl_truncate") or text_cfg["output_dim"]
for split, df in dfs.items():
texts = df["catalog_content"].tolist()
resumable_batch_encode(
items=texts,
encode_batch_fn=lambda batch: text_encoder.encode(batch).numpy(),
output_dim=text_output_dim,
output_path=str(embeddings_dir / f"{split}_text.npy"),
batch_size=text_cfg["batch_size"],
desc=f"{split} text",
)
free_encoder(text_encoder)
logger.info("=== Phase 1/2 complete, text encoder freed from memory ===")
# ---- Phase 2: image embeddings (train + test), then free the encoder ----
logger.info("=== Phase 2/2: image embeddings ===")
image_cfg = config["encoders"]["image"]
image_encoder = build_image_encoder(image_cfg)
image_encoder.ensure_loaded() # detect the real output_dim now, not mid-batch
image_output_dim = image_encoder.output_dim
if image_output_dim != image_cfg.get("output_dim"):
logger.warning(
"configs/base.yaml declares encoders.image.output_dim=%s but the model "
"actually outputs %d — using %d. Update the config to match so training "
"(stage 03) builds the projection layer with the correct input size.",
image_cfg.get("output_dim"), image_output_dim, image_output_dim,
)
def _encode_image_batch(paths):
images = [Image.open(p).convert("RGB") for p in paths]
return image_encoder.encode(images).numpy()
for split, df in dfs.items():
image_paths = df["image_path"].tolist()
resumable_batch_encode(
items=image_paths,
encode_batch_fn=_encode_image_batch,
output_dim=image_output_dim,
output_path=str(embeddings_dir / f"{split}_image.npy"),
batch_size=image_cfg["batch_size"],
desc=f"{split} image",
)
free_encoder(image_encoder)
logger.info("=== Phase 2/2 complete, image encoder freed from memory ===")
# ---- Final sanity check: every split's text/image/row counts line up ----
for split, df in dfs.items():
text_matrix = np.load(embeddings_dir / f"{split}_text.npy")
image_matrix = np.load(embeddings_dir / f"{split}_image.npy")
if text_matrix.shape[0] != len(df) or image_matrix.shape[0] != len(df):
raise PricePredictorError(
f"Embedding count mismatch for split={split}: "
f"text={text_matrix.shape[0]} image={image_matrix.shape[0]} rows={len(df)}"
)
logger.info(
"Verified split=%s: text %s, image %s", split, text_matrix.shape, image_matrix.shape,
)
logger.info("All embeddings extracted and verified.")
def main() -> None:
parser = argparse.ArgumentParser(description="Stage 02: extract and cache embeddings")
parser.add_argument("--config", default="configs/base.yaml")
args = parser.parse_args()
try:
run(args.config)
except PricePredictorError as e:
logger.error("Embedding extraction failed: %s", e)
logger.error("Progress up to the last checkpoint was saved — re-run this exact command to resume.")
sys.exit(1)
except Exception as e:
logger.exception("Unexpected error during embedding extraction: %s", e)
logger.error("Progress up to the last checkpoint was saved — re-run this exact command to resume.")
sys.exit(1)
if __name__ == "__main__":
main() |