update

app.py

@@ -11,15 +11,24 @@ mock_words = [
     "cat", "dog", "rabbit", "hamster"         # Pets
 ]
 
-# Define available models
+# Define available models and load them
 models = {
     'DistilBERT': 'distilbert-base-uncased',
     'BERT': 'bert-base-uncased',
     'RoBERTa': 'roberta-base'
 }
 
+@st.cache(allow_output_mutation=True)
+def load_models():
+    pipelines = {}
+    for name, model_name in models.items():
+        pipelines[name] = pipeline('feature-extraction', model=model_name)
+    return pipelines
+
+pipelines = load_models()
+
 def embed_words(words, model_name):
-    embedder = pipeline('feature-extraction', model=model_name)
+    embedder = pipelines[model_name]
     embeddings = embedder(words)
     return np.array([np.mean(embedding[0], axis=0) for embedding in embeddings])
 
@@ -51,10 +60,9 @@ def main():
 
     # 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, selected_model)
+        clusters = cluster_words(mock_words, model_name)
         display_clusters(clusters)
 
 if __name__ == "__main__":