refactor

load_card_data.py

@@ -1,10 +1,8 @@
 import logging
 import os
-import platform
 from datetime import datetime
-from typing import List, Literal, Optional, Tuple
+from typing import List, Optional, Tuple
 
-import chromadb
 import polars as pl
 import requests
 import stamina
@@ -12,6 +10,7 @@ from chromadb.utils import embedding_functions
 from dotenv import load_dotenv
 from huggingface_hub import InferenceClient
 from tqdm.contrib.concurrent import thread_map
+from utils import get_chroma_client, get_collection
 
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 # Set up logging
@@ -27,30 +26,17 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 EMBEDDING_MODEL_NAME = "Alibaba-NLP/gte-large-en-v1.5"
 EMBEDDING_MODEL_REVISION = "104333d6af6f97649377c2afbde10a7704870c7b"
 INFERENCE_MODEL_URL = (
-    "https://spwy1g6626yhjhpr.us-east-1.aws.endpoints.huggingface.cloud"
+    "https://spwy1g6626yhjhjhpr.us-east-1.aws.endpoints.huggingface.cloud"
 )
 DATASET_PARQUET_URL = (
     "hf://datasets/librarian-bots/dataset_cards_with_metadata/data/train-*.parquet"
 )
 COLLECTION_NAME = "dataset_cards"
-MAX_EMBEDDING_LENGTH = 8192
-
-
-def get_save_path() -> Literal["chroma/"] | Literal["/data/chroma/"]:
-    path = "chroma/" if platform.system() == "Darwin" else "/data/chroma/"
-    logger.info(f"Using save path: {path}")
-    return path
-
-
-SAVE_PATH = get_save_path()
-
 
-def get_chroma_client():
-    logger.info("Initializing Chroma client")
-    return chromadb.PersistentClient(path=SAVE_PATH)
+MAX_EMBEDDING_LENGTH = 8192
 
 
-def get_embedding_function():
+def card_embedding_function():
     logger.info(f"Initializing embedding function with model: {EMBEDDING_MODEL_NAME}")
     return embedding_functions.SentenceTransformerEmbeddingFunction(
         model_name=EMBEDDING_MODEL_NAME,
@@ -59,16 +45,6 @@ def get_embedding_function():
     )
 
 
-def get_collection(chroma_client, embedding_function):
-    logger.info(f"Getting or creating collection: {COLLECTION_NAME}")
-    return chroma_client.create_collection(
-        name=COLLECTION_NAME,
-        get_or_create=True,
-        embedding_function=embedding_function,
-        metadata={"hnsw:space": "cosine"},
-    )
-
-
 def get_last_modified_in_collection(collection) -> datetime | None:
     logger.info("Fetching last modified date from collection")
     try:
@@ -188,9 +164,8 @@ def get_inference_client():
 def refresh_card_data(min_len: int = 250, min_likes: Optional[int] = None):
     logger.info(f"Starting data refresh with min_len={min_len}, min_likes={min_likes}")
     chroma_client = get_chroma_client()
-    embedding_function = get_embedding_function()
-    collection = get_collection(chroma_client, embedding_function)
-
+    embedding_function = card_embedding_function()
+    collection = get_collection(chroma_client, embedding_function, COLLECTION_NAME)
     most_recent = get_last_modified_in_collection(collection)
 
     if data := load_cards(