romybeaute commited on
Commit
b1b8e6c
·
1 Parent(s): 5f5e2fd

Fix zero-shot hang/crash: don't pickle embedding model via cache_data

Browse files
Files changed (1) hide show
  1. app2.py +6 -1
app2.py CHANGED
@@ -828,7 +828,12 @@ def run_zeroshot(_docs, _embeddings, categories, min_similarity, embedding_model
828
  verbose=False,
829
  )
830
  topics, _ = topic_model.fit_transform(list(_docs), np.asarray(_embeddings))
831
- return topics, topic_model.get_topic_info(), topic_model
 
 
 
 
 
832
 
833
 
834
  # =====================================================================
 
828
  verbose=False,
829
  )
830
  topics, _ = topic_model.fit_transform(list(_docs), np.asarray(_embeddings))
831
+ # NOTE: do NOT return topic_model here. @st.cache_data pickles the entire
832
+ # return value, and topic_model holds a reference to the (possibly multi-GB)
833
+ # embedding model — serializing that hangs for minutes and can OOM/crash the
834
+ # Space, dumping the user back to the first tab. Both call sites discard the
835
+ # third value, so return None for it.
836
+ return topics, topic_model.get_topic_info(), None
837
 
838
 
839
  # =====================================================================