Clustering

app.py

@@ -7,6 +7,21 @@ import numpy as np
 import seaborn as sns
 import matplotlib.pyplot as plt
 
+def cluster_examples(messages, embed, nc=3):
+    km = KMeans(
+    n_clusters=nc, init='random',
+    n_init=10, max_iter=300, 
+    tol=1e-04, random_state=0
+    )
+    km = km.fit_predict(embed)
+    for n in range(nc):
+        idxs = [i for i in range(len(km)) if km[i] == n]
+        ms = [messages[i] for i in idxs]
+        st.markdown ("CLUSTER : %d"%n)
+        for m in ms:
+            st.markdown (m)
+
+
 def plot_heatmap(labels, heatmap, rotation=90):
   sns.set(font_scale=1.2)
   fig, ax = plt.subplots()
@@ -29,6 +44,8 @@ st.markdown("This demo uses the sentence_transformers library to plot sentence s
 # Streamlit text boxes
 text = st.text_area('Enter sentences:', value="The sun is hotter than the moon.\nThe sun is very bright.\nI hear that the universe is very large.\nToday is Tuesday.")
 
+nc = st.slider('Select a number of clusters', min_value=1, max_value=15, value=3)
+
 # Model setup
 model = SentenceTransformer('paraphrase-distilroberta-base-v1')
 nltk.download('punkt')
@@ -43,5 +60,6 @@ if text:
         for j,ea in enumerate(embed):
             sim[i][j] = 1.0-cosine(em,ea)
     plot_heatmap(sentences, sim)
+    cluster_examples(sentences, embed, nc)