File size: 2,988 Bytes
6b1bc4b | 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 | """BGE-small (384-d) embeddings of all posts, built EXACTLY like the UniHive `main` backend's personalization worker:
text = '\n'.join(title, body, summary, description)[:12000] (posts have no summary/description -> two empty lines)
vector = fastembed TextEmbedding('BAAI/bge-small-en-v1.5').embed(text) (passage mode, no query prefix), L2-normalised
Two versions per post (same keys as the earlier npz files):
text : title + body only = what the app stores and the engine sees at serving time
fused : title + body + caption in the 'summary' slot = used for TRAINING only (391 posts have an image caption)
No community, no flair in either text. Output: out/post_emb_bge384.npz"""
import json, glob, os, collections
import numpy as np, pandas as pd
from fastembed import TextEmbedding
S = r"C:\Users\tirth\Downloads\Desktop\DATA_FOR_UNIHIVE\enriched"
HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
posts = [json.loads(l) for l in open(os.path.join(S, "posts.jsonl"), encoding="utf8")]
_df = pd.read_csv(os.path.join(HERE, "out", "interactions_v2.csv"))
_have = {p["post_id"] for p in posts}
for r in _df.drop_duplicates("post_id").itertuples(): # posts seen in events but missing from posts.jsonl
if r.post_id not in _have:
posts.append({"post_id": r.post_id, "title": r.title if isinstance(r.title, str) else "", "body": r.body if isinstance(r.body, str) else ""})
cap = collections.defaultdict(list)
for fn in sorted(glob.glob(os.path.join(HERE, "captions*.jsonl"))):
for r in map(json.loads, open(fn, encoding="utf8")):
if r.get("caption"): cap[r["post_id"]].append(r)
s = lambda x: x if isinstance(x, str) else ""
join = lambda title, body, summary: "\n".join([s(title), s(body), s(summary), ""])[:12000] # main: '\n'.join(title, body, summary, description)[:12000]
ids, txt, fus, has_cap = [], [], [], []
for p in posts:
c = " ".join(f"{(r['caption'] or '').strip()} {(r.get('ocr') or '').strip()}".strip() for r in cap.get(p["post_id"], []))
ids.append(p["post_id"]); txt.append(join(p.get("title"), p.get("body"), "")); fus.append(join(p.get("title"), p.get("body"), c)); has_cap.append(bool(c))
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", cache_dir=os.path.join(HERE, "bge_cache"), threads=4)
def enc(texts):
v = np.asarray(list(model.embed(texts, batch_size=32)), dtype=np.float64)
v = v / np.linalg.norm(v, axis=1, keepdims=True); assert v.shape[1] == 384 and np.isfinite(v).all(); return v.astype("float32")
E_t, E_f = enc(txt), enc(fus)
out = os.path.join(HERE, "out", "post_emb_bge384.npz")
np.savez_compressed(out, post_ids=np.array(ids), text=E_t, fused=E_f, has_caption=np.array(has_cap))
cos = (E_t * E_f).sum(1)
print("posts", len(ids), "dim", E_t.shape[1], "| with caption", int(sum(has_cap)), "| cos(text, fused): captioned posts mean %.3f, others mean %.3f" % (cos[np.array(has_cap)].mean(), cos[~np.array(has_cap)].mean()))
|