embed.py

# embed.py
# Purpose: turn text-based columns (e.g., city, notes) or entire customer row into dense embeddings using
# a sentence-transformers model from Hugging Face.


import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
import joblib


MODEL_NAME = 'all-MiniLM-L6-v2' # lightweight sentence-transformer




def build_text_for_embedding(df: pd.DataFrame, text_cols=None):
# combine useful text columns into one string per customer
if text_cols is None:
# try to pick common textual columns
cand = [c for c in df.columns if df[c].dtype == 'object']
text_cols = cand[:3] # at most 3
texts = (df[text_cols].fillna('').astype(str).agg(' | '.join, axis=1))
return texts.tolist()




def embed_texts(texts, model_name=MODEL_NAME, device='cpu'):
model = SentenceTransformer(model_name)
embs = model.encode(texts, show_progress_bar=True)
return embs




if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--features', default='data/features.parquet')
parser.add_argument('--out_emb', default='data/embeddings.npy')
parser.add_argument('--text_cols', nargs='*', default=None)
args = parser.parse_args()


df = pd.read_parquet(args.features)
texts = build_text_for_embedding(df, text_cols=args.text_cols)
embs = embed_texts(texts)
np.save(args.out_emb, embs)
print('Saved embeddings to', args.out_emb)