Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201from argparse import ArgumentParserimport jsonimport osfrom typing import List, Union
from atproto import AsyncIdResolverimport faissfrom fastapi import FastAPI, HTTPExceptionfrom fastapi.responses import HTMLResponsefrom jinja2 import Templateimport numpy as npimport httpxfrom sklearn.cluster import MiniBatchKMeansimport torch
from scripts.utils import get_loggerfrom 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 coolwith 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: ignorelogger.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)