Spaces:
Sleeping
Sleeping
File size: 2,496 Bytes
c3c36d1 |
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 |
"""
Run this LOCALLY to build assets/ from your real dataset.
1) Put your CSV/Parquet with at least 'tagline' (and optional 'description') columns.
2) Adjust INPUT_PATH below.
3) python prepare_assets.py
Then commit assets/ into your Space repo (or upload to a Dataset repo).
"""
import os, json, numpy as np, pandas as pd
from sentence_transformers import SentenceTransformer
import faiss
from logic.cleaning import clean_dataframe
# ---- CHANGE THIS ----
INPUT_PATH = "/mnt/data/hf-slogan-space/data/raw_slogans.csv" # e.g., export from your notebook
ASSETS_DIR = "assets"
MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
NORMALIZE = True # set False if you prefer L2
def main():
os.makedirs(ASSETS_DIR, exist_ok=True)
# Load
if INPUT_PATH.endswith(".csv"):
df = pd.read_csv(INPUT_PATH)
elif INPUT_PATH.endswith(".parquet"):
df = pd.read_parquet(INPUT_PATH)
else:
raise ValueError("Use CSV or Parquet for INPUT_PATH")
# Clean using your real rules
df_clean = clean_dataframe(df)
df_clean.to_parquet(os.path.join(ASSETS_DIR, "slogans_clean.parquet"), index=False)
# Choose text field
if "description" in df_clean.columns:
texts = df_clean["description"].fillna(df_clean["tagline"]).astype(str).tolist()
text_col, fallback_col = "description", "tagline"
else:
texts = df_clean["tagline"].astype(str).tolist()
text_col, fallback_col = "tagline", "tagline"
# Encode
encoder = SentenceTransformer(MODEL_NAME)
emb = encoder.encode(texts, batch_size=64, convert_to_numpy=True, normalize_embeddings=NORMALIZE)
# Save embeddings numpy (optional; not required at runtime)
np.save(os.path.join(ASSETS_DIR, "embeddings.npy"), emb)
# Build FAISS index
dim = emb.shape[1]
if NORMALIZE:
index = faiss.IndexFlatIP(dim) # cosine if normalized
else:
index = faiss.IndexFlatL2(dim)
index.add(emb)
faiss.write_index(index, os.path.join(ASSETS_DIR, "faiss.index"))
meta = {
"model_name": MODEL_NAME,
"dim": int(dim),
"normalized": NORMALIZE,
"metric": "ip" if NORMALIZE else "l2",
"row_count": int(len(df_clean)),
"text_col": text_col,
"fallback_col": fallback_col,
}
with open(os.path.join(ASSETS_DIR, "meta.json"), "w") as f:
json.dump(meta, f, indent=2)
print("✅ Assets built in", ASSETS_DIR)
print(meta)
if __name__ == "__main__":
main()
|