Spaces:
Running on Zero
Running on Zero
File size: 8,568 Bytes
422d4ca | 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 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | """Metadata encoder: turn each row's product-side fields into a fixed-length vector.
Output layout (concatenated):
[ TF-IDF on features_text (META_FEATURE_TFIDF_DIM) # product attributes
| TF-IDF on categories_text (META_CATEGORY_TFIDF_DIM) # product taxonomy
| numeric vector (META_NUM_DIM) # price + ratings + flags
]
Total raw dim = META_TFIDF_DIM + META_NUM_DIM.
The MLP inside the model maps this to META_HIDDEN_DIM. We keep encoding as a
separate, plain-numpy step so it can be fit once on train and reused without
touching the model.
"""
import logging
import pickle
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Sequence
import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import StandardScaler
from . import config as cfg
logger = logging.getLogger(__name__)
# Columns we expect on the input df (produced by preprocess.join_and_clean).
# Missing columns are handled gracefully.
FEATURES_TEXT_COL = "features_text"
CATEGORIES_TEXT_COL = "categories_text"
PRICE_COL = "price"
AVG_RATING_COL = "average_rating"
RATING_NUMBER_COL = "rating_number"
def _safe_text(series: pd.Series) -> pd.Series:
return series.fillna("").astype(str)
def _safe_numeric(series: pd.Series) -> pd.Series:
return pd.to_numeric(series, errors="coerce")
@dataclass
class MetaEncoder:
"""Fit on train df, then transform any df (train/val/test/new) to a (N, D) matrix."""
feature_tfidf_dim: int = cfg.META_FEATURE_TFIDF_DIM
category_tfidf_dim: int = cfg.META_CATEGORY_TFIDF_DIM
num_dim: int = cfg.META_NUM_DIM # price, avg_rating, log_rating_number, price_missing_flag
# Filled by fit()
feature_tfidf: Optional[TfidfVectorizer] = None
category_tfidf: Optional[TfidfVectorizer] = None
scaler: Optional[StandardScaler] = None
price_median_: Optional[float] = None
avg_rating_median_: Optional[float] = None
rating_number_median_: Optional[float] = None
@property
def tfidf_dim(self) -> int:
return self.feature_tfidf_dim + self.category_tfidf_dim
@property
def total_dim(self) -> int:
return self.tfidf_dim + self.num_dim
# ------------------------------------------------------------------ fit
def fit(self, df: pd.DataFrame) -> "MetaEncoder":
# 1. Fit text metadata as separate semantic sources. Keeping features
# and categories apart gives cross-attention tokens real meaning.
feat_text = _safe_text(df.get(FEATURES_TEXT_COL, pd.Series([""] * len(df))))
cat_text = _safe_text(df.get(CATEGORIES_TEXT_COL, pd.Series([""] * len(df))))
self.feature_tfidf = TfidfVectorizer(
max_features=self.feature_tfidf_dim,
ngram_range=(1, 2),
min_df=2,
max_df=0.95,
stop_words="english",
sublinear_tf=True,
)
# If the corpus is too small to have any vocab, TF-IDF will raise.
# Fall back to a single zero-dim by fitting on a synthetic vocab.
try:
self.feature_tfidf.fit(feat_text.str.strip())
except ValueError:
logger.warning("Meta TF-IDF found no usable vocabulary; using zero features.")
self.feature_tfidf = TfidfVectorizer(max_features=1)
self.feature_tfidf.fit(["placeholder"])
self.category_tfidf = TfidfVectorizer(
max_features=self.category_tfidf_dim,
ngram_range=(1, 2),
min_df=2,
max_df=0.98,
stop_words="english",
sublinear_tf=True,
)
try:
self.category_tfidf.fit(cat_text.str.strip())
except ValueError:
logger.warning("Category TF-IDF found no usable vocabulary; using zero features.")
self.category_tfidf = TfidfVectorizer(max_features=1)
self.category_tfidf.fit(["placeholder"])
# 2. Numeric features: price, avg_rating, log(rating_number+1), price_missing_flag
self.price_median_ = float(_safe_numeric(df.get(PRICE_COL, pd.Series([np.nan]))).median())
if not np.isfinite(self.price_median_):
self.price_median_ = 0.0
self.avg_rating_median_ = float(_safe_numeric(df.get(AVG_RATING_COL, pd.Series([np.nan]))).median())
if not np.isfinite(self.avg_rating_median_):
self.avg_rating_median_ = 4.0
self.rating_number_median_ = float(_safe_numeric(df.get(RATING_NUMBER_COL, pd.Series([np.nan]))).median())
if not np.isfinite(self.rating_number_median_):
self.rating_number_median_ = 0.0
num_mat = self._build_numeric(df)
self.scaler = StandardScaler()
self.scaler.fit(num_mat)
return self
# -------------------------------------------------------------- transform
def transform(self, df: pd.DataFrame) -> np.ndarray:
if self.feature_tfidf is None or self.scaler is None:
raise RuntimeError("MetaEncoder.fit() must be called before transform().")
structured = self.transform_structured(df)
return np.hstack([
structured["features"],
structured["categories"],
structured["numeric"],
]).astype(np.float32)
def transform_structured(self, df: pd.DataFrame) -> dict:
"""Return source-separated metadata matrices.
Keys map directly to semantic meta tokens used by the upgraded model:
features -> product attribute text, categories -> taxonomy text,
numeric -> price/rating signals.
"""
if self.feature_tfidf is None or self.category_tfidf is None or self.scaler is None:
raise RuntimeError("MetaEncoder.fit() must be called before transform_structured().")
feat_text = _safe_text(df.get(FEATURES_TEXT_COL, pd.Series([""] * len(df)))).str.strip()
cat_text = _safe_text(df.get(CATEGORIES_TEXT_COL, pd.Series([""] * len(df)))).str.strip()
feat_mat = self.feature_tfidf.transform(feat_text).toarray().astype(np.float32)
cat_mat = self.category_tfidf.transform(cat_text).toarray().astype(np.float32)
feat_mat = self._pad_or_truncate(feat_mat, self.feature_tfidf_dim)
cat_mat = self._pad_or_truncate(cat_mat, self.category_tfidf_dim)
num_mat = self.scaler.transform(self._build_numeric(df)).astype(np.float32)
return {"features": feat_mat, "categories": cat_mat, "numeric": num_mat}
@staticmethod
def _pad_or_truncate(mat: np.ndarray, dim: int) -> np.ndarray:
if mat.shape[1] < dim:
pad = np.zeros((mat.shape[0], dim - mat.shape[1]), dtype=np.float32)
return np.hstack([mat, pad])
if mat.shape[1] > dim:
return mat[:, :dim]
return mat
# ------------------------------------------------------------------ helpers
def _build_numeric(self, df: pd.DataFrame) -> np.ndarray:
price = _safe_numeric(df.get(PRICE_COL, pd.Series([np.nan] * len(df))))
price_missing = price.isna().astype(np.float32).values
price = price.fillna(self.price_median_).astype(np.float32).values
avg = _safe_numeric(df.get(AVG_RATING_COL, pd.Series([np.nan] * len(df))))
avg = avg.fillna(self.avg_rating_median_).astype(np.float32).values
rnum = _safe_numeric(df.get(RATING_NUMBER_COL, pd.Series([np.nan] * len(df))))
rnum = rnum.fillna(self.rating_number_median_).astype(np.float32).values
log_rnum = np.log1p(np.maximum(rnum, 0.0))
return np.vstack([price, avg, log_rnum, price_missing]).T # (N, 4)
# ------------------------------------------------------------------ io
def save(self, path: Path = None) -> Path:
if path is None:
path = cfg.META_ENCODER_PATH
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "wb") as f:
pickle.dump(self, f)
logger.info("Saved MetaEncoder to %s", path)
return path
@classmethod
def load(cls, path: Path = None) -> "MetaEncoder":
if path is None:
path = cfg.META_ENCODER_PATH
with open(path, "rb") as f:
obj = pickle.load(f)
if not isinstance(obj, cls):
raise TypeError(f"Loaded object is not a MetaEncoder: {type(obj)}")
return obj
def fit_and_save(train_df: pd.DataFrame, path: Path = None) -> MetaEncoder:
enc = MetaEncoder().fit(train_df)
enc.save(path)
return enc
|