add other embedding models

app.py

@@ -11,28 +11,51 @@ mock_words = [
     "cat", "dog", "rabbit", "hamster"         # Pets
 ]
 
-# Embedding model
-embedder = pipeline('feature-extraction', model='distilbert-base-uncased')
+# Define available models
+models = {
+    'DistilBERT': 'distilbert-base-uncased',
+    'BERT': 'bert-base-uncased',
+    'RoBERTa': 'roberta-base'
+}
 
-def embed_words(words):
+def embed_words(words, model_name):
+    embedder = pipeline('feature-extraction', model=model_name)
     embeddings = embedder(words)
     return np.array([np.mean(embedding[0], axis=0) for embedding in embeddings])
 
-def cluster_words(words):
-    embeddings = embed_words(words)
+def cluster_words(words, model_name):
+    embeddings = embed_words(words, model_name)
     kmeans = KMeans(n_clusters=4, random_state=0).fit(embeddings)
     clusters = {i: [] for i in range(4)}
     for word, label in zip(words, kmeans.labels_):
         clusters[label].append(word)
+    
+    # Ensure each group has exactly 4 words
+    for i in range(4):
+        if len(clusters[i]) > 4:
+            st.warning(f"Group {i+1} has more than 4 words. Adjusting to show only 4 words.")
+            clusters[i] = clusters[i][:4]
+        elif len(clusters[i]) < 4:
+            st.warning(f"Group {i+1} has less than 4 words. Adjusting by adding placeholder words.")
+            clusters[i].extend(['N/A'] * (4 - len(clusters[i])))
+    
     return clusters
 
+def display_clusters(clusters):
+    for i, words in clusters.items():
+        st.markdown(f"### Group {i+1}")
+        st.write(", ".join(words))
+
 def main():
     st.title("NYT Connections Solver")
-    
+
+    # Dropdown menu for selecting the embedding model
+    model_name = st.selectbox("Select Embedding Model", list(models.keys()))
+    selected_model = models[model_name]
+
     if st.button("Generate Clusters"):
-        clusters = cluster_words(mock_words)
-        for i, words in clusters.items():
-            st.write(f"Group {i+1}: {', '.join(words)}")
+        clusters = cluster_words(mock_words, selected_model)
+        display_clusters(clusters)
 
 if __name__ == "__main__":
     main()