fix decorator and words

app.py

@@ -5,10 +5,10 @@ import numpy as np
 
 # Mock data
 mock_words = [
-    "apple", "banana", "cherry", "date",      # Fruits
-    "car", "truck", "bus", "bicycle",         # Vehicles
-    "red", "blue", "green", "yellow",         # Colors
-    "cat", "dog", "rabbit", "hamster"         # Pets
+    "apple", "banana", "cherry", "date",  # Fruits
+    "car", "truck", "bus", "bicycle",  # Vehicles
+    "red", "blue", "green", "yellow",  # Colors
+    "cat", "dog", "rabbit", "hamster"  # Pets
 ]
 
 # Define available models and load them
@@ -18,7 +18,7 @@ models = {
     'RoBERTa': 'roberta-base'
 }
 
-@st.cache(allow_output_mutation=True)
+@st.cache_resource
 def load_models():
     pipelines = {}
     for name, model_name in models.items():
@@ -28,6 +28,9 @@ def load_models():
 pipelines = load_models()
 
 def embed_words(words, model_name):
+    """
+    Embed the given words using the specified model and return the averaged embeddings.
+    """
     embedder = pipelines[model_name]
     embeddings = embedder(words)
     return np.array([np.mean(embedding[0], axis=0) for embedding in embeddings])
@@ -37,17 +40,8 @@ def cluster_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])))
-    
+        if len(clusters[label]) < 4:
+            clusters[label].append(word)
     return clusters
 
 def display_clusters(clusters):
@@ -57,13 +51,16 @@ def display_clusters(clusters):
 
 def main():
     st.title("NYT Connections Solver")
-
+    st.write("This app demonstrates solving the NYT Connections game using word embeddings and clustering.")
+    st.write("Select an embedding model from the dropdown menu and click 'Generate Clusters' to see the grouped words.")
+    
     # Dropdown menu for selecting the embedding model
     model_name = st.selectbox("Select Embedding Model", list(models.keys()))
 
     if st.button("Generate Clusters"):
-        clusters = cluster_words(mock_words, model_name)
+        with st.spinner("Generating clusters..."):
+            clusters = cluster_words(mock_words, model_name)
         display_clusters(clusters)
 
 if __name__ == "__main__":
-    main()
+    main()