refactor: update model and embedding configurations, enhance logging for database setup

main.py

@@ -24,8 +24,8 @@ load_dotenv(override=True)
 HF_TOKEN = os.getenv("HF_TOKEN")
 login(token=HF_TOKEN)
 # Configuration constants
-MODEL_NAME = "davanstrien/SmolLM2-360M-tldr-sft-2025-02-12_15-13"
-EMBEDDING_MODEL = "nomic-ai/modernbert-embed-base"
+MODEL_NAME = "davanstrien/Smol-Hub-tldr"
+EMBEDDING_MODEL = "Qwen/Qwen3-Embedding-0.6B"
 BATCH_SIZE = 2000
 CACHE_TTL = "24h"
 TRENDING_CACHE_TTL = "1h"  # 15 minutes cache for trending data
@@ -38,9 +38,7 @@ else:
     DEVICE = "cpu"
 
 
-tokenizer = AutoTokenizer.from_pretrained(
-    "davanstrien/SmolLM2-360M-tldr-sft-2025-02-12_15-13"
-)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"  # turn on HF_TRANSFER
 # Set up logging
@@ -90,7 +88,7 @@ app.add_middleware(
 def get_embedding_function():
     logger.info(f"Using device: {DEVICE}")
     return embedding_functions.SentenceTransformerEmbeddingFunction(
-        model_name="nomic-ai/modernbert-embed-base", device=DEVICE
+        model_name="Qwen/Qwen3-Embedding-0.6B", device=DEVICE
     )
 
 
@@ -135,24 +133,64 @@ def setup_database():
             logger.info(f"Most recent record in DB from: {latest_update}")
             logger.info(f"Oldest record in DB from: {min(last_modifieds)}")
 
+            # Log sample of existing timestamps for debugging
+            sample_timestamps = sorted(last_modifieds, reverse=True)[:5]
+            logger.info(f"Sample of most recent DB timestamps: {sample_timestamps}")
+
         # Filter and process only newer records
         df = df.select(["datasetId", "summary", "likes", "downloads", "last_modified"])
 
-        # Log some stats about the incoming data
-        sample_dates = df.select("last_modified").limit(5).collect()
-        logger.info(f"Sample of incoming dates: {sample_dates}")
-
+        # Log some stats about the incoming data BEFORE collecting
         total_incoming = df.select(pl.len()).collect().item()
-        logger.info(f"Total incoming records: {total_incoming}")
+        logger.info(f"Total incoming records from source: {total_incoming}")
+
+        # Get sample of dates to understand the data
+        sample_df = (
+            df.select(["datasetId", "last_modified"])
+            .sort("last_modified", descending=True)
+            .limit(10)
+            .collect()
+        )
+        logger.info("Sample of most recent incoming records:")
+        for row in sample_df.iter_rows():
+            logger.info(f"  {row[0]}: {row[1]}")
 
         if latest_update:
             logger.info(f"Filtering records newer than {latest_update}")
+            logger.info(f"Latest update type: {type(latest_update)}")
+
+            # Get date range before filtering
+            date_stats = df.select(
+                [
+                    pl.col("last_modified").min().alias("min_date"),
+                    pl.col("last_modified").max().alias("max_date"),
+                ]
+            ).collect()
+            logger.info(f"Incoming data date range: {date_stats.row(0)}")
+
             # Ensure last_modified is datetime before comparison
             df = df.with_columns(pl.col("last_modified").str.to_datetime())
             df = df.filter(pl.col("last_modified") > latest_update)
             filtered_count = df.select(pl.len()).collect().item()
             logger.info(f"Found {filtered_count} records to update after filtering")
 
+            if filtered_count == 0:
+                logger.warning(
+                    "No new records found after filtering! This might indicate a problem."
+                )
+                # Log a few records that were just below the cutoff
+                just_before = (
+                    df.select(["datasetId", "last_modified"])
+                    .filter(pl.col("last_modified") <= latest_update)
+                    .sort("last_modified", descending=True)
+                    .limit(5)
+                    .collect()
+                )
+                if len(just_before) > 0:
+                    logger.info("Records just before cutoff:")
+                    for row in just_before.iter_rows():
+                        logger.info(f"  {row[0]}: {row[1]}")
+
         df = df.collect()
         total_rows = len(df)
 
@@ -170,8 +208,26 @@ def setup_database():
                     f"({batch_df['last_modified'].min()} to {batch_df['last_modified'].max()})"
                 )
 
+                ids_to_upsert = batch_df.select(["datasetId"]).to_series().to_list()
+
+                # Log first few IDs being upserted
+                logger.info(f"Upserting IDs (first 5): {ids_to_upsert[:5]}")
+
+                # Check if any of these already exist
+                existing_check = dataset_collection.get(
+                    ids=ids_to_upsert[:5], include=["metadatas"]
+                )
+                if existing_check["ids"]:
+                    logger.info(
+                        f"Found {len(existing_check['ids'])} existing records in this batch sample"
+                    )
+                    for idx, id_ in enumerate(existing_check["ids"]):
+                        logger.info(
+                            f"  Existing: {id_} - last_modified: {existing_check['metadatas'][idx].get('last_modified')}"
+                        )
+
                 dataset_collection.upsert(
-                    ids=batch_df.select(["datasetId"]).to_series().to_list(),
+                    ids=ids_to_upsert,
                     documents=batch_df.select(["summary"]).to_series().to_list(),
                     metadatas=[
                         {
@@ -188,18 +244,55 @@ def setup_database():
                 )
                 logger.info(f"Processed {i + batch_size:,} / {total_rows:,} records")
 
-        logger.info(
-            f"Database initialized with {dataset_collection.count():,} total rows"
-        )
+        # Final validation
+        final_count = dataset_collection.count()
+        logger.info(f"Database initialized with {final_count:,} total rows")
+
+        # Verify the update worked by checking latest records
+        if final_count > 0:
+            final_metadata = dataset_collection.get(include=["metadatas"], limit=5)
+            final_timestamps = [
+                dateutil.parser.parse(m.get("last_modified"))
+                for m in final_metadata.get("metadatas")
+            ]
+            if final_timestamps:
+                latest_after_update = max(final_timestamps)
+                logger.info(f"Latest record after update: {latest_after_update}")
+                if latest_update and latest_after_update <= latest_update:
+                    logger.error(
+                        "WARNING: No new records were added! Latest timestamp hasn't changed."
+                    )
+                elif latest_update:
+                    logger.info(
+                        f"Successfully added records from {latest_update} to {latest_after_update}"
+                    )
 
         # Load model data
         model_lazy_df = pl.scan_parquet(
             "hf://datasets/davanstrien/models_with_metadata_and_summaries/data/train-*.parquet"
         )
         model_row_count = model_lazy_df.select(pl.len()).collect().item()
-        logger.info(f"Row count of new model data: {model_row_count}")
+        logger.info(f"Total model records in source: {model_row_count}")
+
+        # Get the most recent last_modified date from the model collection
+        model_latest_update = None
+        if model_collection.count() > 0:
+            model_metadata = model_collection.get(include=["metadatas"]).get(
+                "metadatas"
+            )
+            logger.info(
+                f"Found {len(model_metadata)} existing model records in collection"
+            )
+
+            model_last_modifieds = [
+                dateutil.parser.parse(m.get("last_modified")) for m in model_metadata
+            ]
+            model_latest_update = max(model_last_modifieds)
+            logger.info(f"Most recent model record in DB from: {model_latest_update}")
 
-        if model_collection.count() < model_row_count:
+        # Always process models to handle updates (not just new additions)
+        should_update_models = True
+        if model_latest_update:
             schema = model_lazy_df.collect_schema()
             select_columns = [
                 "modelId",