add device detection for model inference and improve dataset collection logging

main.py

@@ -12,6 +12,7 @@ from contextlib import asynccontextmanager
 import polars as pl
 from huggingface_hub import HfApi
 from transformers import AutoTokenizer
+import torch
 
 # Configuration constants
 MODEL_NAME = "davanstrien/SmolLM2-360M-tldr-sft-2025-02-12_15-13"
@@ -19,6 +20,13 @@ EMBEDDING_MODEL = "nomic-ai/modernbert-embed-base"
 BATCH_SIZE = 1000
 CACHE_TTL = "60"
 
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+elif torch.backends.mps.is_available():
+    DEVICE = "mps"
+else:
+    DEVICE = "cpu"
+
 hf_api = HfApi()
 
 
@@ -72,8 +80,9 @@ app.add_middleware(
 
 # Define the embedding function at module level
 def get_embedding_function():
+    logger.info(f"Using device: {DEVICE}")
     return embedding_functions.SentenceTransformerEmbeddingFunction(
-        model_name="nomic-ai/modernbert-embed-base"
+        model_name="nomic-ai/modernbert-embed-base", device=DEVICE
     )
 
 
@@ -95,7 +104,7 @@ def setup_database():
             metadata={"hnsw:space": "cosine"},
         )
 
-        # TODO incremental updates
+        # Load dataset data
         df = pl.scan_parquet(
             "hf://datasets/davanstrien/datasets_with_metadata_and_summaries/data/train-*.parquet"
         )
@@ -103,14 +112,21 @@ def setup_database():
             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:
+        logger.info(f"Row count of dataset data: {row_count}")
+
+        # Check if we need to update the collection
+        current_count = dataset_collection.count()
+        logger.info(f"Current dataset collection count: {current_count}")
+
+        if current_count < row_count:
+            logger.info(
+                f"Updating dataset collection with {row_count - current_count} new records"
+            )
             # 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):
@@ -148,7 +164,6 @@ def setup_database():
                 ["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):