Update app.py

app.py

@@ -10,9 +10,6 @@ import faiss
 import gradio as gr
 from sklearn.metrics.pairwise import cosine_similarity
 from huggingface_hub import hf_hub_download, login
-from pyspark.sql import SparkSession
-from pyspark.sql.functions import col, udf, monotonically_increasing_id, collect_list, concat_ws
-from pyspark.sql.types import StringType
 from huggingface_hub import HfApi
 
 # Load token from Hugging Face Secrets 
@@ -34,12 +31,15 @@ def load_reddit_split(subreddit_name):
 # Combine subreddit data
 combined_dataset = chain(*(load_reddit_split(sub) for sub in target_subreddits))
 
-if "JAVA_HOME" not in os.environ:
-    os.environ["JAVA_HOME"] = "/usr/lib/jvm/java-11-openjdk-amd64"  # Common path on HF Spaces/Debian
+import pandas as pd
+import re
+from itertools import islice
+
+# Load a sample of the dataset (e.g., 100,000 records for performance)
+comments = [{"body": ex["body"]} for ex in islice(combined_dataset, 100000)]
 
-# PySpark session
-spark = SparkSession.builder.getOrCreate()
-df = spark.createDataFrame([{"body": ex["body"]} for ex in islice(combined_dataset, 100000)])
+# Convert to DataFrame
+df = pd.DataFrame(comments)
 
 # Clean text function
 def clean_body(text):
@@ -48,17 +48,17 @@ def clean_body(text):
     text = re.sub(r"[^a-zA-Z\s]", "", text)
     return re.sub(r"\s+", " ", text).strip()
 
-clean_udf = udf(clean_body, StringType())
-df_clean = df.withColumn("clean", clean_udf(col("body")))
+# Apply cleaning
+df["clean"] = df["body"].apply(clean_body)
 
-# Chunking
+# Chunk every 5 rows
 chunk_size = 5
-df_indexed = df_clean.withColumn("row_num", monotonically_increasing_id())
-df_indexed = df_indexed.withColumn("chunk_id", (col("row_num") / chunk_size).cast("int"))
-df_chunked = df_indexed.groupBy("chunk_id").agg(concat_ws(" ", collect_list("clean")).alias("chunk_text"))
+df["chunk_id"] = df.index // chunk_size
+df_chunked = df.groupby("chunk_id")["clean"].apply(lambda texts: " ".join(texts)).reset_index()
+df_chunked.rename(columns={"clean": "chunk_text"}, inplace=True)
 
-# Collect for embedding
-chunked_comments = df_chunked.select("chunk_text").rdd.map(lambda x: x[0]).collect()
+# 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))