chroma and models

main.py

@@ -10,26 +10,14 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 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 huggingface_hub import HfApi
 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
+
+# Configuration constants
+MODEL_NAME = "davanstrien/SmolLM2-360M-tldr-sft-2025-02-12_15-13"
+EMBEDDING_MODEL = "nomic-ai/modernbert-embed-base"
+BATCH_SIZE = 1000
+CACHE_TTL = "30"
 
 hf_api = HfApi()
 
@@ -74,7 +62,7 @@ app.add_middleware(
     allow_origins=[
         "https://*.hf.space",  # Allow all Hugging Face Spaces
         "https://*.huggingface.co",  # Allow all Hugging Face domains
-        # "http://localhost:5500",  # Allow localhost:5500 # TODO remove before prod
+        "http://localhost:5500",  # Allow localhost:5500 # TODO remove before prod
     ],
     allow_credentials=True,
     allow_methods=["*"],
@@ -93,12 +81,20 @@ def setup_database():
     try:
         embedding_function = get_embedding_function()
 
-        # Create collection with embedding function
+        # Create dataset collection
         dataset_collection = client.get_or_create_collection(
             embedding_function=embedding_function,
             name="dataset_cards",
             metadata={"hnsw:space": "cosine"},
         )
+
+        # Create model collection
+        model_collection = client.get_or_create_collection(
+            embedding_function=embedding_function,
+            name="model_cards",
+            metadata={"hnsw:space": "cosine"},
+        )
+
         # TODO incremental updates
         df = pl.scan_parquet(
             "hf://datasets/davanstrien/datasets_with_metadata_and_summaries/data/train-*.parquet"
@@ -139,42 +135,48 @@ def setup_database():
                 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")
+
+        # Load model data
+        model_df = pl.scan_parquet(
+            "hf://datasets/davanstrien/models_with_metadata_and_summaries/data/train-*.parquet"
+        )
+        model_row_count = model_df.select(pl.len()).collect().item()
+        logger.info(f"Row count of new model data: {model_row_count}")
+
+        if model_collection.count() < model_row_count:
+            model_df = model_df.select(
+                ["modelId", "summary", "likes", "downloads", "last_modified"]
+            )
+            model_df = model_df.collect()
+            BATCH_SIZE = 1000
+            total_rows = len(model_df)
+
+            for i in range(0, total_rows, BATCH_SIZE):
+                batch_df = model_df.slice(i, min(BATCH_SIZE, total_rows - i))
+
+                model_collection.upsert(
+                    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),
+                            "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:,} model rows"
+                )
+
+        logger.info(
+            f"Model database initialized with {model_collection.count():,} rows"
+        )
 
     except Exception as e:
         logger.error(f"Setup error: {e}")
@@ -196,6 +198,18 @@ class QueryResponse(BaseModel):
     results: List[QueryResult]
 
 
+class ModelQueryResult(BaseModel):
+    model_id: str
+    similarity: float
+    summary: str
+    likes: int
+    downloads: int
+
+
+class ModelQueryResponse(BaseModel):
+    results: List[ModelQueryResult]
+
+
 @app.get("/")
 async def redirect_to_docs():
     from fastapi.responses import RedirectResponse
@@ -204,7 +218,7 @@ async def redirect_to_docs():
 
 
 @app.get("/search/datasets", response_model=QueryResponse)
-@cache(ttl="10m")
+@cache(ttl=CACHE_TTL)
 async def search_datasets(
     query: str,
     k: int = Query(default=5, ge=1, le=100),
@@ -235,22 +249,7 @@ async def search_datasets(
         )
 
         # 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"],
-                )
-            )
-
-        # 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]
+        query_results = process_search_results(results, "dataset", k, sort_by)
 
         return QueryResponse(results=query_results)
 
@@ -260,7 +259,7 @@ async def search_datasets(
 
 
 @app.get("/similarity/datasets", response_model=QueryResponse)
-@cache(ttl="10m")
+@cache(ttl=CACHE_TTL)
 async def find_similar_datasets(
     dataset_id: str,
     k: int = Query(default=5, ge=1, le=100),
@@ -298,25 +297,9 @@ async def find_similar_datasets(
         )
 
         # 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]
+        query_results = process_search_results(
+            results, "dataset", k, sort_by, dataset_id
+        )
 
         return QueryResponse(results=query_results)
 
@@ -327,6 +310,119 @@ async def find_similar_datasets(
         raise HTTPException(status_code=500, detail="Similarity search failed")
 
 
+@app.get("/search/models", response_model=ModelQueryResponse)
+@cache(ttl=CACHE_TTL)
+async def search_models(
+    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:
+        collection = client.get_collection(
+            name="model_cards", embedding_function=get_embedding_function()
+        )
+
+        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,
+        )
+
+        query_results = process_search_results(results, "model", k, sort_by)
+
+        return ModelQueryResponse(results=query_results)
+
+    except Exception as e:
+        logger.error(f"Model search error: {str(e)}")
+        raise HTTPException(status_code=500, detail="Model search failed")
+
+
+@app.get("/similarity/models", response_model=ModelQueryResponse)
+@cache(ttl=CACHE_TTL)
+async def find_similar_models(
+    model_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:
+        collection = client.get_collection("model_cards")
+
+        results = collection.get(ids=[model_id], include=["embeddings"])
+
+        if not results["ids"]:
+            raise HTTPException(
+                status_code=404, detail=f"Model ID '{model_id}' not found"
+            )
+
+        results = collection.query(
+            query_embeddings=[results["embeddings"][0]],
+            n_results=k * 4 if sort_by != "similarity" else k + 1,
+            where={
+                "$and": [
+                    {"likes": {"$gte": min_likes}},
+                    {"downloads": {"$gte": min_downloads}},
+                ]
+            }
+            if min_likes > 0 or min_downloads > 0
+            else None,
+        )
+
+        query_results = process_search_results(results, "model", k, sort_by, model_id)
+
+        return ModelQueryResponse(results=query_results)
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Model similarity search error: {str(e)}")
+        raise HTTPException(status_code=500, detail="Model similarity search failed")
+
+
+def process_search_results(results, id_field, k, sort_by, exclude_id=None):
+    """Process search results into a standardized format."""
+    query_results = []
+    for i in range(len(results["ids"][0])):
+        current_id = results["ids"][0][i]
+        if exclude_id and current_id == exclude_id:
+            continue
+
+        result = {
+            f"{id_field}_id": current_id,
+            "similarity": float(results["distances"][0][i]),
+            "summary": results["documents"][0][i],
+            "likes": results["metadatas"][0][i]["likes"],
+            "downloads": results["metadatas"][0][i]["downloads"],
+        }
+
+        if id_field == "dataset":
+            query_results.append(QueryResult(**result))
+        else:
+            query_results.append(ModelQueryResult(**result))
+
+    if sort_by != "similarity":
+        query_results.sort(key=lambda x: getattr(x, sort_by), reverse=True)
+        query_results = query_results[:k]
+    elif exclude_id:  # We fetched extra for similarity + exclude_id case
+        query_results = query_results[:k]
+
+    return query_results
+
+
 if __name__ == "__main__":
     import uvicorn