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