Create ingest/cluster.py

ingest/cluster.py

@@ -1,29 +1,21 @@
-from sentence_transformers import SentenceTransformer
-import faiss, numpy as np
-import plotly.graph_objects as go
+import faiss
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
 
-model = SentenceTransformer("all-MiniLM-L6-v2")
+def semantic_clusters(documents, k=5):
+    """
+    documents: list[str]
+    returns: list[int] cluster ids
+    """
+    if len(documents) < 2:
+        return [0] * len(documents)
 
-def semantic_cluster_plot(results):
-    texts = [r["title"] + " " + r["snippet"] for r in results]
-    if not texts:
-        return go.Figure()
+    vectorizer = TfidfVectorizer(max_features=512, stop_words="english")
+    vectors = vectorizer.fit_transform(documents).toarray().astype("float32")
 
-    embeddings = model.encode(texts)
-    dim = embeddings.shape[1]
+    dim = vectors.shape[1]
     index = faiss.IndexFlatL2(dim)
-    index.add(np.array(embeddings))
+    index.add(vectors)
 
-    # simple 2D projection (first 2 dims for HF safety)
-    x, y = embeddings[:,0], embeddings[:,1]
-
-    fig = go.Figure(
-        data=go.Scatter(
-            x=x, y=y,
-            mode="markers",
-            text=[r["title"] for r in results],
-            marker=dict(size=8)
-        )
-    )
-    fig.update_layout(title="Semantic Document Clusters")
-    return fig
+    _, labels = index.search(vectors, 1)
+    return labels.flatten().tolist()