Adding filter for ids

src/utilities.py

@@ -2,6 +2,7 @@ import os
 
 import numpy as np
 import pandas as pd
+import requests
 from datasets import Dataset, DownloadMode, load_dataset
 from gradio_client import Client
 
@@ -12,6 +13,7 @@ USERNAME = os.environ["USERNAME"]
 OG_DATASET = f"{USERNAME}/dataset-creator-reddit-{SUBREDDIT}"
 PROCESSED_DATASET = os.environ['PROCESSED_DATASET']
 embeddings_space = f"{USERNAME}/nomic-embeddings"
+FILTER_IDS_URL = "https://huggingface.co/spaces/reddit-tools-HF/dataset-creator-reddit-bestofredditorupdates/raw/main/filter_ids.json"
 
 logger = setup_logger(__name__)
 
@@ -36,6 +38,9 @@ def merge_and_update_datasets(dataset, original_dataset):
     odf = original_dataset['train'].to_pandas()
     df = dataset['train'].to_pandas()
 
+    # Filter ODF in-case we missed any
+    odf = remove_filtered_rows(odf, FILTER_IDS_URL)
+
     # Step 1: Merge df onto odf
     # We'll bring in 'content' and 'embedding' from df to compare and possibly update 'embedding'
     merged_df = pd.merge(odf, df[['id', 'content', 'embedding']], on='id', how='left', suffixes=('_odf', ''))
@@ -60,6 +65,29 @@ def merge_and_update_datasets(dataset, original_dataset):
     return dataset, updated_row_count
 
 
+def remove_filtered_rows(df: pd.DataFrame, url: str) -> pd.DataFrame:
+    """
+    Removes rows from the DataFrame where the 'id' is present in the JSON file at the given URL.
+
+    :param df: Input DataFrame to be filtered.
+    :param url: URL to the JSON file containing the filter IDs.
+    :return: DataFrame with rows containing IDs present in the JSON file removed.
+    """
+
+    # Load filter IDs from JSON file at the URL
+    response = requests.get(url)
+    filter_ids = response.json()
+
+    logger.info(f"Loaded {len(filter_ids)} IDs from {url}")
+
+    # Remove the rows with IDs present in filter_ids
+    filtered_df = df[~df['id'].astype(str).isin(filter_ids)]
+
+    logger.info(f"Filtered {len(df) - len(filtered_df)} rows from the DataFrame")
+
+    return filtered_df
+
+
 def update_embeddings(content, client):
     embedding = client.predict('search_document: ' + content, api_name="/embed")
     return np.array(embedding)