Spaces:
Build error
Build error
| """ | |
| Feature store: pre-computes and caches all embeddings at startup. | |
| Retrieves embeddings by entity ID in O(1) from in-memory dict, | |
| with SQLite as persistent backup. | |
| """ | |
| import numpy as np | |
| from typing import Dict | |
| from artwork_bandit.db import database | |
| class FeatureStore: | |
| def __init__(self, db, nlp_encoder, vision_encoder): | |
| self.db = db | |
| self.nlp = nlp_encoder | |
| self.vis = vision_encoder | |
| self.user_emb: Dict[str, np.ndarray] = {} | |
| self.content_emb: Dict[str, np.ndarray] = {} | |
| self.artwork_emb: Dict[str, np.ndarray] = {} | |
| self.artworks_by_content: Dict[str, list] = {} | |
| def precompute_all(self, users, contents, artworks): | |
| # users | |
| print(f"[PRECOMPUTE] Encoding {len(users)} users...") | |
| for i, u in enumerate(users): | |
| if i % 50 == 0: | |
| print(f"[PRECOMPUTE] User {i}/{len(users)}") | |
| try: | |
| emb = self.nlp.encode_user(u) | |
| self.user_emb[u['user_id']] = emb | |
| try: | |
| database.save_embedding('user', u['user_id'], emb) | |
| except Exception: | |
| pass | |
| except Exception as e: | |
| print(f"[PRECOMPUTE] Error encoding user {u.get('user_id')}: {e}") | |
| print(f"[PRECOMPUTE] Completed {len(users)} users") | |
| # contents | |
| print(f"[PRECOMPUTE] Encoding {len(contents)} contents...") | |
| for i, c in enumerate(contents): | |
| if i % 10 == 0: | |
| print(f"[PRECOMPUTE] Content {i}/{len(contents)}") | |
| try: | |
| emb = self.nlp.encode_content(c) | |
| self.content_emb[c['content_id']] = emb | |
| try: | |
| database.save_embedding('content', c['content_id'], emb) | |
| except Exception: | |
| pass | |
| except Exception as e: | |
| print(f"[PRECOMPUTE] Error encoding content {c.get('content_id')}: {e}") | |
| print(f"[PRECOMPUTE] Completed {len(contents)} contents") | |
| # artworks | |
| print(f"[PRECOMPUTE] Encoding {len(artworks)} artworks...") | |
| for i, a in enumerate(artworks): | |
| if i % 50 == 0: | |
| print(f"[PRECOMPUTE] Artwork {i}/{len(artworks)}") | |
| try: | |
| emb = self.vis.encode_artwork(a) | |
| self.artwork_emb[a['artwork_id']] = emb | |
| try: | |
| database.save_embedding('artwork', a['artwork_id'], emb) | |
| except Exception: | |
| pass | |
| self.artworks_by_content.setdefault(a['content_id'], []).append(a['artwork_id']) | |
| except Exception as e: | |
| print(f"[PRECOMPUTE] Error encoding artwork {a.get('artwork_id')}: {e}") | |
| print(f"[PRECOMPUTE] Completed {len(artworks)} artworks") | |
| print("[PRECOMPUTE] All embeddings precomputed successfully") | |
| def get_user_embedding(self, user_id: str): | |
| return self.user_emb.get(user_id) | |
| def get_content_embedding(self, content_id: str): | |
| return self.content_emb.get(content_id) | |
| def get_artwork_embedding(self, artwork_id: str): | |
| return self.artwork_emb.get(artwork_id) | |
| def build_context_vector(self, user_id, content_id, artwork_id): | |
| u = self.get_user_embedding(user_id) | |
| c = self.get_content_embedding(content_id) | |
| a = self.get_artwork_embedding(artwork_id) | |
| if u is None or c is None or a is None: | |
| return None | |
| # ensure vector order and concatenation [user(384)|content(384)|artwork(512)] | |
| return np.concatenate([u, c, a]) | |
| def get_artworks_for_content(self, content_id: str): | |
| return self.artworks_by_content.get(content_id, []) | |