Update app.py

app.py

@@ -15,30 +15,24 @@ from huggingface_hub import HfApi
 # Load token from Hugging Face Secrets 
 HF_TOKEN = os.environ.get("RedditSemanticSearch")
 
-# Function to stream JSONL Reddit files from HF Hub
 from datasets import load_dataset
+from itertools import islice
 
 # Define target subreddits
 target_subreddits = ["askscience", "gaming", "technology", "todayilearned", "programming"]
 
-# Load full Reddit dataset (assumes it's pre-split by subreddit or has a field)
-dataset_splits = [load_dataset("HuggingFaceGECLM/REDDIT_comments", split=sub, streaming=True) for sub in target_subreddits]
-
-# Filter only relevant subreddits
-dataset = dataset.filter(lambda x: x["subreddit"] in target_subreddits)
-
-# Take a sample (to limit memory for now)
-comments = [{"body": ex["body"]} for ex in dataset.select(range(100000))]
+# Load streaming dataset for each subreddit and combine
+def stream_subreddit_data(subreddit):
+    return load_dataset("HuggingFaceGECLM/REDDIT_comments", split=subreddit, streaming=True)
 
-import pandas as pd
-import re
-from itertools import islice
+# Combine streams
+combined_dataset = chain(*(stream_subreddit_data(sub) for sub in target_subreddits))
 
-# Load a sample of the dataset (e.g., 100,000 records for performance)
-comments = [{"body": ex["body"]} for ex in islice(combined_dataset, 100000)]
+# Sample up to 100,000 comments from the combined stream
+comments = list(islice(combined_dataset, 100000))
 
-# Convert to DataFrame
-df = pd.DataFrame(comments)
+# Extract text and subreddit
+df = pd.DataFrame([{"body": ex["body"], "subreddit": ex["subreddit"]} for ex in comments])
 
 # Clean text function
 def clean_body(text):
@@ -59,13 +53,8 @@ df_chunked.rename(columns={"clean": "chunk_text"}, inplace=True)
 # Final list for embedding
 chunked_comments = df_chunked["chunk_text"].tolist()
 
-# Create subreddit labels
-combined_dataset = chain(*(load_reddit_split(sub) for sub in target_subreddits))
-subreddit_labels = []
-for example in combined_dataset:
-    subreddit_labels.append(example["subreddit_name_prefixed"])
-    if len(subreddit_labels) >= len(chunked_comments):
-        break
+# Create subreddit labels (reused from original list)
+subreddit_labels = df["subreddit"].tolist()[:len(chunked_comments)]
 
 # Tokenize
 def clean_text(text):