from argparse import ArgumentParser import json import os from typing import List, Union from atproto import AsyncIdResolver import faiss from fastapi import FastAPI, HTTPException from fastapi.responses import HTMLResponse from jinja2 import Template import numpy as np import httpx from sklearn.cluster import MiniBatchKMeans import torch from scripts.utils import get_logger from scripts.training.models import FollowEmbedModule logger = get_logger(__name__) BSKY_URL = "https://public.api.bsky.app" for var in ["CHECKPOINT_PATH", "DID_ID_MAP_PATH", "NUM_CLUSTERS"]: if var not in os.environ: raise ValueError(f"Missing required env var: {var}") checkpoint_path = os.environ["CHECKPOINT_PATH"] did_id_map_path = os.environ["DID_ID_MAP_PATH"] num_clusters = int(os.environ["NUM_CLUSTERS"]) logger.info("Loading model from checkpoint...") model = FollowEmbedModule.load_from_checkpoint(checkpoint_path, device="cpu") source_embed = model.source_embed.embed.weight.cpu().detach().to(torch.float16).numpy() target_embed = model.target_embed.embed.weight.cpu().detach().to(torch.float16).numpy() del model # Manual memory mgmt is cool with open(did_id_map_path, "r") as in_file: did_id_map: dict[str, int] = json.load(in_file) id_did_map = {value: key for key, value in did_id_map.items()} logger.info("Building index...") target_index = faiss.IndexHNSWFlat(target_embed.shape[1], 32) target_index.add(target_embed) # type: ignore logger.info("Index constructed") logger.info("Initializing cluster model...") kmeans = MiniBatchKMeans( n_clusters=num_clusters, random_state=123, batch_size=1024, ) logger.info("Training k-means model...") kmeans.fit(target_embed) logger.info("Cluster model training complete") async def get_follows(did: str) -> List[str]: follows = [] try: async with httpx.AsyncClient() as client: data = await client.get( BSKY_URL + "/xrpc/app.bsky.graph.getFollows", params={"actor": did, "limit": 100}, headers={"Content-Type": "application/json"}, ) data = data.json() for follow in data["follows"]: follows.append(follow["did"]) except Exception as e: raise e if "cursor" not in data: return follows # Limit to 2000 follows per account while data["cursor"] and len(follows) < 2000: try: async with httpx.AsyncClient() as client: data = await client.get( BSKY_URL + "/xrpc/app.bsky.graph.getFollows", params={"actor": did, "limit": 100, "cursor": data["cursor"]}, headers={"Content-Type": "application/json"}, ) data = data.json() for follow in data["follows"]: follows.append(follow["did"]) if "cursor" not in data: break except Exception as e: raise e return follows async def get_profiles(dids: list[str]): # logger.info("Getting data on:\n" + "\n".join(dids)) profiles = [] try: async with httpx.AsyncClient() as client: data = await client.get( BSKY_URL + "/xrpc/app.bsky.actor.getProfiles", params={"actors": dids}, headers={"Content-Type": "application/json"}, ) data = data.json() for profile in data["profiles"]: profiles.append(profile) return profiles except Exception as e: raise e async def get_recs(did: str): follows = await get_follows(did) logger.info(f"Number of follows: {len(follows)}") user_id = did_id_map[did] user_vec = np.expand_dims(source_embed[user_id, :], axis=0) # (1,256) _, neighbors = target_index.search(user_vec, 25) # type: ignore neigh_dids: list[str] = [] for neigh in neighbors[0]: neigh_did: str = id_did_map[neigh.item()] if neigh_did not in follows: neigh_dids.append(neigh_did) profiles = await get_profiles(neigh_dids) return profiles async def get_cluster_recs(did: str): follows = await get_follows(did) follow_ids = [] for follow_did in follows: if follow_did in did_id_map: follow_ids.append(did_id_map[follow_did]) # For each follow, check centroid assignment followed_user_vectors = target_embed[follow_ids, :] centroids = kmeans.predict(followed_user_vectors) # Calculate percent weighting for top 5 centroids centroids, counts = np.unique(centroids, return_counts=True) sort_indexes = counts.argsort() centroids = centroids[sort_indexes[::-1]][:5] counts = counts[sort_indexes[::-1]][:5] weights = counts / counts.sum() # Get neighbors of centroids based on weighted relevance neighbor_dids: list[str] = [] samples = np.random.multinomial(25, weights) for idx, sample_count in enumerate(samples): # logger.info(f"Drawing {sample_count} samples from centroid: {centroids[idx]}") if sample_count < 1: continue centroid_vec = np.expand_dims( kmeans.cluster_centers_[centroids[idx], :], axis=0 # type: ignore ) _, neighbors = target_index.search(centroid_vec, sample_count.item()) # type: ignore for neigh in neighbors[0]: neigh_did: str = id_did_map[neigh.item()] if neigh_did not in follows: # logger.info(f"Recommending did: {neigh_did}") neighbor_dids.append(neigh_did) profiles = await get_profiles(neighbor_dids) return profiles app = FastAPI() with open( os.path.join(os.path.dirname(os.path.abspath(__file__)), "page.html"), "r" ) as template_file: template_str = template_file.read() template = Template(template_str) @app.get("/") async def render_page(query: Union[str, None] = None): logger.info(f"Received query: {query}") if not query: result = template.render(profiles=[], query="") return HTMLResponse(content=result) # If it doesn't start with 'did' resolve as handle if query[:4] == "did": did = query else: resolver = AsyncIdResolver() did = await resolver.handle.resolve(query) if did is None: return HTTPException(status_code=500, detail="Couldn't resolve handle") recs = await get_cluster_recs(did) for rec in recs: if "avatar" in rec: rec["avatar"] = rec["avatar"].replace("avatar", "avatar_thumbnail") result = template.render(profiles=recs, query=query) return HTMLResponse(content=result)