Merge branch 'agent-paper-main' into joanne-tutorialagent

1 files changed, 32 insertions, 9 deletions

diff --git a/src/server/chunker/pdf_chunker.py b/src/server/chunker/pdf_chunker.py
index e34753176..04d9f51a4 100644
--- a/src/server/chunker/pdf_chunker.py
+++ b/src/server/chunker/pdf_chunker.py
@@ -701,20 +701,43 @@ class Document:
 
         :return: The generated summary of the document.
         """
-        num_clusters = min(10, len(self.chunks))  # Set number of clusters for KMeans, capped at 10
-        kmeans = KMeans(n_clusters=num_clusters, random_state=42)  # Initialize KMeans with 10 clusters
-        doc_chunks = [chunk['values'] for chunk in self.chunks if 'values' in chunk]  # Extract embeddings
-        cluster_labels = kmeans.fit_predict(doc_chunks)  # Assign each chunk to a cluster
+        # num_clusters = min(10, len(self.chunks))  # Set number of clusters for KMeans, capped at 10
+        # kmeans = KMeans(n_clusters=num_clusters, random_state=42)  # Initialize KMeans with 10 clusters
+        # doc_chunks = [chunk['values'] for chunk in self.chunks if 'values' in chunk]  # Extract embeddings
+        # cluster_labels = kmeans.fit_predict(doc_chunks)  # Assign each chunk to a cluster
+
+        doc_chunks = [chunk['values'] for chunk in self.chunks if 'values' in chunk]
+        if not doc_chunks:
+            raise ValueError("No valid embedded chunks to summarize.")
+
+        # Remove duplicates (e.g., from OCR-ed blank pages or repeated captions)
+        unique_chunks = np.unique(np.array(doc_chunks), axis=0)
+
+        # Dynamically scale number of clusters to available signal
+        num_clusters = min(10, len(unique_chunks))
+        kmeans = KMeans(n_clusters=num_clusters, random_state=42).fit(unique_chunks)
+
+        # Predict cluster labels for original chunks (not just unique ones)
+        cluster_labels = kmeans.predict(np.array(doc_chunks))
+
 
         # Select representative chunks from each cluster
         selected_chunks = []
         for i in range(num_clusters):
-            cluster_chunks = [chunk for chunk, label in zip(self.chunks, cluster_labels) if label == i]  # Get all chunks in this cluster
-            cluster_embs = [emb for emb, label in zip(doc_chunks, cluster_labels) if label == i]  # Get embeddings for this cluster
+            # cluster_chunks = [chunk for chunk, label in zip(self.chunks, cluster_labels) if label == i]  # Get all chunks in this cluster
+            # cluster_embs = [emb for emb, label in zip(doc_chunks, cluster_labels) if label == i]  # Get embeddings for this cluster
+
+            cluster_idxs = np.where(cluster_labels == i)[0]
+            if len(cluster_idxs) == 0:
+                continue  # skip empty clusters (shouldn't happen after downsizing)
+
             centroid = kmeans.cluster_centers_[i]  # Get the centroid of the cluster
-            distances = [np.linalg.norm(np.array(emb) - centroid) for emb in cluster_embs]  # Compute distance to centroid
-            closest_chunk = cluster_chunks[np.argmin(distances)]  # Select chunk closest to the centroid
-            selected_chunks.append(closest_chunk)
+            distances = [np.linalg.norm(doc_chunks[idx] - centroid) for idx in cluster_idxs]
+            closest_idx = cluster_idxs[int(np.argmin(distances))]
+            selected_chunks.append(self.chunks[closest_idx])
+            # distances = [np.linalg.norm(np.array(emb) - centroid) for emb in cluster_embs]  # Compute distance to centroid
+            # closest_chunk = cluster_chunks[np.argmin(distances)]  # Select chunk closest to the centroid
+            # selected_chunks.append(closest_chunk)
 
         # Combine selected chunks into a summary
         combined_text = "\n\n".join([chunk['metadata']['text'] for chunk in selected_chunks])  # Concatenate chunk texts