switch to chromadb

main.py

@@ -2,13 +2,41 @@ import logging
 import os
 from typing import List
 import sys
-import duckdb
-from cashews import cache  # Add this import
+import chromadb
+from chromadb.utils import embedding_functions
+from cashews import cache
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from sentence_transformers import SentenceTransformer
 from contextlib import asynccontextmanager
+import polars as pl
+from huggingface_hub import hf_hub_url, DatasetCard, ModelCard, HfApi
+from datetime import datetime, timedelta
+from typing import Generator
+from huggingface_hub import ModelInfo, DatasetInfo
+import stamina
+import logging
+import polars as pl
+from huggingface_hub import dataset_info
+from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer
+import stamina
+from tqdm.contrib.concurrent import thread_map
+from datasets import Dataset, Value, Sequence
+import datasets
+import os
+from dotenv import load_dotenv
+from huggingface_hub import get_inference_endpoint
+from huggingface_hub import AsyncInferenceClient
+import asyncio
+from typing import List
+
+hf_api = HfApi()
+
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "davanstrien/SmolLM2-360M-tldr-sft-2025-02-12_15-13"
+)
 
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"  # turn on HF_TRANSFER
 # Set up logging
@@ -22,15 +50,20 @@ DATA_DIR = "data" if LOCAL else "/data"
 # Configure cache
 cache.setup("mem://", size_limit="4gb")
 
+# Initialize ChromaDB client
+client = chromadb.PersistentClient(path=f"{DATA_DIR}/chroma")
+
 
 # Initialize FastAPI app
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup: nothing special needed here since model and DB are initialized at module level
+    # Setup
+    setup_database()
+
     yield
+
     # Cleanup
     await cache.close()
-    con.close()
 
 
 app = FastAPI(lifespan=lifespan)
@@ -48,62 +81,100 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Initialize model and DuckDB
-model = SentenceTransformer("nomic-ai/modernbert-embed-base", backend="onnx")
-embedding_dim = model.get_sentence_embedding_dimension()
-
-# Database setup with fallback
-db_path = f"{DATA_DIR}/vector_store.db"
-try:
-    # Create directory if it doesn't exist
-    os.makedirs(os.path.dirname(db_path), exist_ok=True)
-    con = duckdb.connect(db_path)
-    logger.info(f"Connected to persistent database at {db_path}")
-except (OSError, PermissionError) as e:
-    logger.warning(
-        f"Could not create/access {db_path}. Falling back to in-memory database. Error: {e}"
-    )
-    con = duckdb.connect(":memory:")
 
-# Initialize VSS extension
-con.sql("INSTALL vss; LOAD vss;")
-con.sql("SET hnsw_enable_experimental_persistence=true;")
+# Define the embedding function at module level
+def get_embedding_function():
+    return embedding_functions.SentenceTransformerEmbeddingFunction(
+        model_name="nomic-ai/modernbert-embed-base"
+    )
 
 
 def setup_database():
     try:
-        # Create table with properly typed embeddings
-        con.sql(f"""
-            CREATE TABLE IF NOT EXISTS model_cards AS 
-            SELECT *, embeddings::FLOAT[{embedding_dim}] as embeddings_float 
-            FROM 'hf://datasets/davanstrien/outputs-embeddings/**/*.parquet';
-        """)
-
-        # Check if index exists
-        index_exists = (
-            con.sql("""
-            SELECT COUNT(*) as count 
-            FROM duckdb_indexes 
-            WHERE index_name = 'my_hnsw_index';
-        """).fetchone()[0]
-            > 0
-        )
-
-        if index_exists:
-            # Drop existing index
-            con.sql("DROP INDEX my_hnsw_index;")
-            logger.info("Dropped existing HNSW index")
-
-        # Create/Recreate HNSW index
-        con.sql("""
-            CREATE INDEX my_hnsw_index ON model_cards 
-            USING HNSW (embeddings_float) WITH (metric = 'cosine');
-        """)
-        logger.info("Created/Recreated HNSW index")
+        embedding_function = get_embedding_function()
 
-        # Log the number of rows in the database
-        row_count = con.sql("SELECT COUNT(*) as count FROM model_cards").fetchone()[0]
-        logger.info(f"Database initialized with {row_count:,} rows")
+        # Create collection with embedding function
+        dataset_collection = client.get_or_create_collection(
+            embedding_function=embedding_function,
+            name="dataset_cards",
+            metadata={"hnsw:space": "cosine"},
+        )
+        # TODO incremental updates
+        df = pl.scan_parquet(
+            "hf://datasets/davanstrien/datasets_with_metadata_and_summaries/data/train-*.parquet"
+        )
+        df = df.filter(
+            pl.col("datasetId").str.contains_any(["open-llm-leaderboard-old/"]).not_()
+        )
+        row_count = df.select(pl.len()).collect().item()
+        logger.info(f"Row count of new data: {row_count}")
+        if dataset_collection.count() < row_count:
+            # Load parquet files and upsert into ChromaDB
+            df = df.select(
+                ["datasetId", "summary", "likes", "downloads", "last_modified"]
+            )
+            df = df.collect()
+            BATCH_SIZE = 1000
+            total_rows = len(df)
+
+            for i in range(0, total_rows, BATCH_SIZE):
+                batch_df = df.slice(i, min(BATCH_SIZE, total_rows - i))
+
+                dataset_collection.upsert(
+                    ids=batch_df.select(["datasetId"]).to_series().to_list(),
+                    documents=batch_df.select(["summary"]).to_series().to_list(),
+                    metadatas=[
+                        {
+                            "likes": int(likes),
+                            "downloads": int(downloads),
+                            "last_modified": str(last_modified),
+                        }
+                        for likes, downloads, last_modified in zip(
+                            batch_df.select(["likes"]).to_series().to_list(),
+                            batch_df.select(["downloads"]).to_series().to_list(),
+                            batch_df.select(["last_modified"]).to_series().to_list(),
+                        )
+                    ],
+                )
+                logger.info(f"Processed {i + len(batch_df):,} / {total_rows:,} rows")
+
+        logger.info(f"Database initialized with {dataset_collection.count():,} rows")
+        # model_collection = client.get_or_create_collection(
+        #     embedding_function=embedding_function,
+        #     name="model_cards",
+        #     metadata={"hnsw:space": "cosine"},
+        # )
+
+        # # If collection is empty, load data from parquet files
+        # if model_collection.count() == 0:
+        #     # Load parquet files and insert into ChromaDB
+        #     df = pl.scan_parquet(
+        #         "hf://datasets/librarian-bots/model_cards_with_metadata/data/train-*.parquet"
+        #     )
+        #     df = df.select(["modelId", "likes", "downloads"])
+        #     df = df.collect()
+        #     df = df.sample(n=1000)  # TODO remove for prod
+        #     # Process in batches of 1000
+        #     BATCH_SIZE = 1000
+        #     total_rows = len(df)
+
+        #     for i in range(0, total_rows, BATCH_SIZE):
+        #         batch_df = df.slice(i, min(BATCH_SIZE, total_rows - i))
+
+        #         model_collection.add(
+        #             ids=batch_df.select(["modelId"]).to_series().to_list(),
+        #             documents=batch_df.select(["summary"]).to_series().to_list(),
+        #             metadatas=[
+        #                 {"likes": int(likes), "downloads": int(downloads)}
+        #                 for likes, downloads in zip(
+        #                     batch_df.select(["likes"]).to_series().to_list(),
+        #                     batch_df.select(["downloads"]).to_series().to_list(),
+        #                 )
+        #             ],
+        #         )
+        #         logger.info(f"Processed {i + len(batch_df):,} / {total_rows:,} rows")
+
+        # logger.info(f"Database initialized with {model_collection.count():,} rows")
 
     except Exception as e:
         logger.error(f"Setup error: {e}")
@@ -134,39 +205,54 @@ async def redirect_to_docs():
 
 @app.get("/search/datasets", response_model=QueryResponse)
 @cache(ttl="10m")
-async def search_datasets(query: str, k: int = Query(default=5, ge=1, le=100)):
+async def search_datasets(
+    query: str,
+    k: int = Query(default=5, ge=1, le=100),
+    sort_by: str = Query(
+        default="similarity", enum=["similarity", "likes", "downloads"]
+    ),
+    min_likes: int = Query(default=0, ge=0),
+    min_downloads: int = Query(default=0, ge=0),
+):
     try:
-        query_embedding = model.encode(f"search_query: {query}").tolist()
-
-        # Updated SQL query to include likes and downloads
-        result = con.sql(f"""
-            SELECT 
-                datasetId as dataset_id,
-                1 - array_cosine_distance(
-                    embeddings_float::FLOAT[{embedding_dim}], 
-                    {query_embedding}::FLOAT[{embedding_dim}]
-                ) as similarity,
-                summary,
-                likes,
-                downloads
-            FROM model_cards
-            ORDER BY similarity DESC
-            LIMIT {k};
-        """).df()
-
-        # Updated result conversion
-        results = [
-            QueryResult(
-                dataset_id=row["dataset_id"],
-                similarity=float(row["similarity"]),
-                summary=row["summary"],
-                likes=int(row["likes"]),
-                downloads=int(row["downloads"]),
+        # Get collection with proper embedding function
+        collection = client.get_collection(
+            name="dataset_cards", embedding_function=get_embedding_function()
+        )
+
+        # Query ChromaDB
+        results = collection.query(
+            query_texts=[f"search_query: {query}"],
+            n_results=k * 4 if sort_by != "similarity" else k,
+            where={
+                "$and": [
+                    {"likes": {"$gte": min_likes}},
+                    {"downloads": {"$gte": min_downloads}},
+                ]
+            }
+            if min_likes > 0 or min_downloads > 0
+            else None,
+        )
+
+        # Process results
+        query_results = []
+        for i in range(len(results["ids"][0])):
+            query_results.append(
+                QueryResult(
+                    dataset_id=results["ids"][0][i],
+                    similarity=float(results["distances"][0][i]),
+                    summary=results["documents"][0][i],
+                    likes=results["metadatas"][0][i]["likes"],
+                    downloads=results["metadatas"][0][i]["downloads"],
+                )
             )
-            for _, row in result.iterrows()
-        ]
 
-        return QueryResponse(results=results)
+        # Sort results if needed
+        if sort_by != "similarity":
+            query_results.sort(key=lambda x: getattr(x, sort_by), reverse=True)
+            query_results = query_results[:k]
+
+        return QueryResponse(results=query_results)
 
     except Exception as e:
         logger.error(f"Search error: {str(e)}")
@@ -176,52 +262,63 @@ async def search_datasets(query: str, k: int = Query(default=5, ge=1, le=100)):
 @app.get("/similarity/datasets", response_model=QueryResponse)
 @cache(ttl="10m")
 async def find_similar_datasets(
-    dataset_id: str, k: int = Query(default=5, ge=1, le=100)
+    dataset_id: str,
+    k: int = Query(default=5, ge=1, le=100),
+    sort_by: str = Query(
+        default="similarity", enum=["similarity", "likes", "downloads"]
+    ),
+    min_likes: int = Query(default=0, ge=0),
+    min_downloads: int = Query(default=0, ge=0),
 ):
     try:
-        # First, get the embedding for the input dataset_id
-        reference_embedding = con.sql(f"""
-            SELECT embeddings_float
-            FROM model_cards
-            WHERE datasetId = '{dataset_id}'
-            LIMIT 1;
-        """).df()
-
-        if reference_embedding.empty:
+        collection = client.get_collection("dataset_cards")
+
+        # Get the reference document
+        results = collection.get(ids=[dataset_id], include=["embeddings"])
+
+        if not results["ids"]:
             raise HTTPException(
                 status_code=404, detail=f"Dataset ID '{dataset_id}' not found"
             )
 
-        # Updated similarity search query to include likes and downloads
-        result = con.sql(f"""
-            SELECT 
-                datasetId as dataset_id,
-                1 - array_cosine_distance(
-                    embeddings_float::FLOAT[{embedding_dim}], 
-                    (SELECT embeddings_float FROM model_cards WHERE datasetId = '{dataset_id}' LIMIT 1)
-                ) as similarity,
-                summary,
-                likes,
-                downloads
-            FROM model_cards
-            WHERE datasetId != '{dataset_id}'
-            ORDER BY similarity DESC
-            LIMIT {k};
-        """).df()
-
-        # Updated result conversion
-        results = [
-            QueryResult(
-                dataset_id=row["dataset_id"],
-                similarity=float(row["similarity"]),
-                summary=row["summary"],
-                likes=int(row["likes"]),
-                downloads=int(row["downloads"]),
-            )
-            for _, row in result.iterrows()
-        ]
+        # Query using the embedding
+        results = collection.query(
+            query_embeddings=[results["embeddings"][0]],
+            n_results=k * 4
+            if sort_by != "similarity"
+            else k + 1,  # +1 to account for self-match
+            where={
+                "$and": [
+                    {"likes": {"$gte": min_likes}},
+                    {"downloads": {"$gte": min_downloads}},
+                ]
+            }
+            if min_likes > 0 or min_downloads > 0
+            else None,
+        )
 
-        return QueryResponse(results=results)
+        # Process results (excluding the query dataset itself)
+        query_results = []
+        for i in range(len(results["ids"][0])):
+            if results["ids"][0][i] != dataset_id:
+                query_results.append(
+                    QueryResult(
+                        dataset_id=results["ids"][0][i],
+                        similarity=float(results["distances"][0][i]),
+                        summary=results["documents"][0][i],
+                        likes=results["metadatas"][0][i]["likes"],
+                        downloads=results["metadatas"][0][i]["downloads"],
+                    )
+                )
+
+        # Sort results if needed
+        if sort_by != "similarity":
+            query_results.sort(key=lambda x: getattr(x, sort_by), reverse=True)
+            query_results = query_results[:k]
+        else:
+            query_results = query_results[:k]
+
+        return QueryResponse(results=query_results)
 
     except HTTPException:
         raise