add param_count to model metadata and update requirements for compatibility

main.py

@@ -4,7 +4,7 @@ import os
 import sys
 from contextlib import asynccontextmanager
 from datetime import datetime
-from typing import List
+from typing import List, Optional
 
 import chromadb
 import dateutil.parser
@@ -194,7 +194,14 @@ def setup_database():
 
         if model_collection.count() < model_row_count:
             model_df = model_df.select(
-                ["modelId", "summary", "likes", "downloads", "last_modified"]
+                [
+                    "modelId",
+                    "summary",
+                    "likes",
+                    "downloads",
+                    "last_modified",
+                    "param_count",
+                ]
             )
             model_df = model_df.collect()
             total_rows = len(model_df)
@@ -210,11 +217,15 @@ def setup_database():
                             "likes": int(likes),
                             "downloads": int(downloads),
                             "last_modified": str(last_modified),
+                            "param_count": int(param_count)
+                            if param_count is not None
+                            else 0,
                         }
-                        for likes, downloads, last_modified in zip(
+                        for likes, downloads, last_modified, param_count 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(),
+                            batch_df.select(["param_count"]).to_series().to_list(),
                         )
                     ],
                 )
@@ -252,6 +263,7 @@ class ModelQueryResult(BaseModel):
     summary: str
     likes: int
     downloads: int
+    param_count: Optional[int] = None
 
 
 class ModelQueryResponse(BaseModel):
@@ -471,6 +483,10 @@ async def process_search_results(results, id_field, k, sort_by, exclude_id=None)
             "downloads": results["metadatas"][0][i]["downloads"],
         }
 
+        # Add param_count for models if it exists in metadata
+        if id_field == "model" and "param_count" in results["metadatas"][0][i]:
+            result["param_count"] = results["metadatas"][0][i]["param_count"]
+
         if id_field == "dataset":
             query_results.append(QueryResult(**result))
         else:
@@ -546,21 +562,24 @@ async def get_trending_models_with_summaries(
 
         # Fetch summaries from ChromaDB
         collection = client.get_collection("model_cards")
-        summaries = collection.get(ids=model_ids, include=["documents"])
+        summaries = collection.get(ids=model_ids, include=["documents", "metadatas"])
 
-        # Create mapping of model_id to summary
+        # Create mapping of model_id to summary and metadata
         id_to_summary = dict(zip(summaries["ids"], summaries["documents"]))
+        id_to_metadata = dict(zip(summaries["ids"], summaries["metadatas"]))
 
         # Combine data
         results = []
         for model in trending_models:
             if model["modelId"] in id_to_summary:
+                metadata = id_to_metadata.get(model["modelId"], {})
                 result = ModelQueryResult(
                     model_id=model["modelId"],
                     similarity=1.0,  # Not applicable for trending
                     summary=id_to_summary[model["modelId"]],
                     likes=model.get("likes", 0),
                     downloads=model.get("downloads", 0),
+                    param_count=metadata.get("param_count"),
                 )
                 results.append(result)
 

requirements.in

@@ -1,6 +1,6 @@
 aiohttp
 cashews
-chromadb
+chromadb==1.0.0b0
 datasets
 einops
 fastapi

requirements.txt

@@ -19,10 +19,13 @@ anyio==4.8.0
 asgiref==3.8.1
     # via opentelemetry-instrumentation-asgi
 attrs==25.1.0
-    # via aiohttp
+    # via
+    #   aiohttp
+    #   jsonschema
+    #   referencing
 backoff==2.2.1
     # via posthog
-bcrypt==4.2.1
+bcrypt==4.3.0
     # via chromadb
 build==1.2.2.post1
     # via chromadb
@@ -40,7 +43,7 @@ charset-normalizer==3.4.1
     # via requests
 chroma-hnswlib==0.7.6
     # via chromadb
-chromadb==0.6.3
+chromadb==1.0.0b0
     # via -r requirements.in
 click==8.1.8
     # via
@@ -59,11 +62,13 @@ dill==0.3.8
     # via
     #   datasets
     #   multiprocess
+distro==1.9.0
+    # via posthog
 durationpy==0.9
     # via kubernetes
 einops==0.8.1
     # via -r requirements.in
-fastapi==0.115.8
+fastapi==0.115.9
     # via
     #   -r requirements.in
     #   chromadb
@@ -135,6 +140,10 @@ jinja2==3.1.5
     # via torch
 joblib==1.4.2
     # via scikit-learn
+jsonschema==4.23.0
+    # via chromadb
+jsonschema-specifications==2024.10.1
+    # via jsonschema
 kubernetes==32.0.1
     # via chromadb
 markdown-it-py==3.0.0
@@ -228,9 +237,9 @@ pandas==2.2.3
     # via datasets
 pillow==11.1.0
     # via sentence-transformers
-polars==1.23.0
+polars==1.24.0
     # via -r requirements.in
-posthog==3.15.1
+posthog==3.18.0
     # via chromadb
 propcache==0.3.0
     # via
@@ -281,6 +290,10 @@ pyyaml==6.0.2
     #   kubernetes
     #   transformers
     #   uvicorn
+referencing==0.36.2
+    # via
+    #   jsonschema
+    #   jsonschema-specifications
 regex==2024.11.6
     # via transformers
 requests==2.32.3
@@ -297,6 +310,10 @@ rich==13.9.4
     # via
     #   chromadb
     #   typer
+rpds-py==0.23.1
+    # via
+    #   jsonschema
+    #   referencing
 rsa==4.9
     # via google-auth
 safetensors==0.5.3
@@ -309,7 +326,7 @@ scipy==1.15.2
     #   sentence-transformers
 sentence-transformers==3.4.1
     # via -r requirements.in
-setuptools==75.8.1
+setuptools==75.8.2
     # via torch
 shellingham==1.5.4
     # via typer
@@ -350,7 +367,7 @@ tqdm==4.67.1
     #   transformers
 transformers==4.49.0
     # via sentence-transformers
-typer==0.15.1
+typer==0.15.2
     # via chromadb
 typing-extensions==4.12.2
     # via
@@ -361,6 +378,7 @@ typing-extensions==4.12.2
     #   opentelemetry-sdk
     #   pydantic
     #   pydantic-core
+    #   referencing
     #   torch
     #   typer
 tzdata==2025.1