File size: 1,672 Bytes
a7c372c | 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 | """Turn a catalog row into the text document the embedding model reads."""
import numpy as np
KEEP_CATEGORIES = {
"Single-player", "Multi-player", "Co-op", "Online Co-op", "LAN Co-op", "Shared/Split Screen Co-op",
"PvP", "Online PvP", "Shared/Split Screen PvP", "MMO", "VR Support", "VR Only", "Cross-Platform Multiplayer",
}
def build_doc(row, rng: np.random.Generator | None = None, drop_tags: bool = False, max_tags: int = 15) -> str:
"""Deterministic when rng is None. With rng, returns a randomly degraded view for contrastive training."""
aug = rng is not None
tags = list(row["tags"])[:max_tags]
if drop_tags or (aug and rng.random() < 0.5):
tags = []
elif aug and tags:
keep = rng.random(len(tags)) >= 0.3
tags = [t for t, k in zip(tags, keep) if k]
genres = list(row["genres"])
cats = [c for c in row["categories"] if c in KEEP_CATEGORIES][:5]
desc = (row["short"] + " " + row["about"]).strip()
if aug and desc:
if rng.random() < 0.3 and (tags or genres):
desc = ""
else:
words = desc.split()
cut = int(len(words) * rng.uniform(0.3, 1.0)) if len(words) > 20 else len(words)
desc = " ".join(words[:cut])
parts = [row["name"] + "."]
if genres:
parts.append("Genres: " + ", ".join(genres) + ".")
if tags:
parts.append("Tags: " + ", ".join(tags) + ".")
if cats:
parts.append("Modes: " + ", ".join(cats) + ".")
if len(row["developers"]):
parts.append("Developer: " + ", ".join(list(row["developers"])[:2]) + ".")
if desc:
parts.append(desc)
return " ".join(parts)
|