add trending sorting option and fetch trending scores for datasets and models

main.py

@@ -1,21 +1,23 @@
+import asyncio
 import logging
 import os
-from typing import List
 import sys
+from contextlib import asynccontextmanager
+from datetime import datetime
+from typing import List
+
 import chromadb
-from chromadb.utils import embedding_functions
+import dateutil.parser
+import httpx
+import polars as pl
+import torch
 from cashews import cache
+from chromadb.utils import embedding_functions
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
+from huggingface_hub import HfApi, model_info
 from pydantic import BaseModel
-from contextlib import asynccontextmanager
-import polars as pl
-from huggingface_hub import HfApi
 from transformers import AutoTokenizer
-import torch
-import dateutil.parser
-import httpx
-from datetime import datetime
 
 # Configuration constants
 MODEL_NAME = "davanstrien/SmolLM2-360M-tldr-sft-2025-02-12_15-13"
@@ -272,18 +274,16 @@ 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"]
+        default="similarity", enum=["similarity", "likes", "downloads", "trending"]
     ),
     min_likes: int = Query(default=0, ge=0),
     min_downloads: int = Query(default=0, ge=0),
 ):
     try:
-        # 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,
@@ -297,8 +297,7 @@ async def search_datasets(
             else None,
         )
 
-        # Process results
-        query_results = process_search_results(results, "dataset", k, sort_by)
+        query_results = await process_search_results(results, "dataset", k, sort_by)
 
         return QueryResponse(results=query_results)
 
@@ -313,7 +312,7 @@ async def find_similar_datasets(
     dataset_id: str,
     k: int = Query(default=5, ge=1, le=100),
     sort_by: str = Query(
-        default="similarity", enum=["similarity", "likes", "downloads"]
+        default="similarity", enum=["similarity", "likes", "downloads", "trending"]
     ),
     min_likes: int = Query(default=0, ge=0),
     min_downloads: int = Query(default=0, ge=0),
@@ -321,7 +320,6 @@ async def find_similar_datasets(
     try:
         collection = client.get_collection("dataset_cards")
 
-        # Get the reference document
         results = collection.get(ids=[dataset_id], include=["embeddings"])
 
         if not results["ids"]:
@@ -329,12 +327,9 @@ async def find_similar_datasets(
                 status_code=404, detail=f"Dataset ID '{dataset_id}' not found"
             )
 
-        # 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
+            n_results=k * 4 if sort_by != "similarity" else k + 1,
             where={
                 "$and": [
                     {"likes": {"$gte": min_likes}},
@@ -345,8 +340,7 @@ async def find_similar_datasets(
             else None,
         )
 
-        # Process results (excluding the query dataset itself)
-        query_results = process_search_results(
+        query_results = await process_search_results(
             results, "dataset", k, sort_by, dataset_id
         )
 
@@ -365,7 +359,7 @@ 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"]
+        default="similarity", enum=["similarity", "likes", "downloads", "trending"]
     ),
     min_likes: int = Query(default=0, ge=0),
     min_downloads: int = Query(default=0, ge=0),
@@ -388,7 +382,7 @@ async def search_models(
             else None,
         )
 
-        query_results = process_search_results(results, "model", k, sort_by)
+        query_results = await process_search_results(results, "model", k, sort_by)
 
         return ModelQueryResponse(results=query_results)
 
@@ -431,7 +425,9 @@ async def find_similar_models(
             else None,
         )
 
-        query_results = process_search_results(results, "model", k, sort_by, model_id)
+        query_results = await process_search_results(
+            results, "model", k, sort_by, model_id
+        )
 
         return ModelQueryResponse(results=query_results)
 
@@ -442,9 +438,29 @@ async def find_similar_models(
         raise HTTPException(status_code=500, detail="Model similarity search failed")
 
 
-def process_search_results(results, id_field, k, sort_by, exclude_id=None):
+@cache(ttl="1h")
+async def get_trending_score(item_id: str, item_type: str) -> float:
+    """Fetch trending score for a model or dataset from HuggingFace API"""
+    try:
+        async with httpx.AsyncClient() as client:
+            endpoint = "models" if item_type == "model" else "datasets"
+            response = await client.get(
+                f"https://huggingface.co/api/{endpoint}/{item_id}?expand=trendingScore"
+            )
+            response.raise_for_status()
+            return response.json().get("trendingScore", 0)
+    except Exception as e:
+        logger.error(
+            f"Error fetching trending score for {item_type} {item_id}: {str(e)}"
+        )
+        return 0
+
+
+async def process_search_results(results, id_field, k, sort_by, exclude_id=None):
     """Process search results into a standardized format."""
     query_results = []
+
+    # Create base results
     for i in range(len(results["ids"][0])):
         current_id = results["ids"][0][i]
         if exclude_id and current_id == exclude_id:
@@ -463,7 +479,31 @@ def process_search_results(results, id_field, k, sort_by, exclude_id=None):
         else:
             query_results.append(ModelQueryResult(**result))
 
-    if sort_by != "similarity":
+    # Handle sorting
+    if sort_by == "trending":
+        # Fetch trending scores for all results
+        trending_scores = {}
+        async with httpx.AsyncClient() as client:
+            tasks = [
+                get_trending_score(
+                    getattr(result, f"{id_field}_id"),
+                    "model" if id_field == "model" else "dataset",
+                )
+                for result in query_results
+            ]
+            scores = await asyncio.gather(*tasks)
+            trending_scores = {
+                getattr(result, f"{id_field}_id"): score
+                for result, score in zip(query_results, scores)
+            }
+
+        # Sort by trending score
+        query_results.sort(
+            key=lambda x: trending_scores.get(getattr(x, f"{id_field}_id"), 0),
+            reverse=True,
+        )
+        query_results = query_results[:k]
+    elif 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