Update model version.

pages/predict.py

@@ -36,16 +36,21 @@ st.markdown(f"**Query:** {st.session_state.query['source_node'].replace('_', ' '
 
 @st.cache_data(show_spinner = 'Downloading AI model...')
 def get_embeddings():
+    # Get checkpoint name
+    # best_ckpt = "2024_05_22_11_59_43_epoch=18-step=22912"
+    best_ckpt = "2024_05_15_13_05_33_epoch=2-step=40383"
+    # best_ckpt = "2024_03_29_04_12_52_epoch=3-step=54291"
+
     # Get paths to embeddings, relation weights, and edge types
     # with st.spinner('Downloading AI model...'):
     embed_path = hf_hub_download(repo_id="ayushnoori/galaxy",
-                                filename="2024_03_29_04_12_52_epoch=3-step=54291_embeddings.pt",
+                                filename=(best_ckpt + "-thresh=4000_embeddings.pt"),
                                 token=st.secrets["HF_TOKEN"])
     relation_weights_path = hf_hub_download(repo_id="ayushnoori/galaxy",
-                                            filename="2024_03_29_04_12_52_epoch=3-step=54291_relation_weights.pt",
+                                            filename=(best_ckpt + "_relation_weights.pt"),
                                             token=st.secrets["HF_TOKEN"])
     edge_types_path = hf_hub_download(repo_id="ayushnoori/galaxy",
-                                        filename="2024_03_29_04_12_52_epoch=3-step=54291_edge_types.pt",
+                                        filename=(best_ckpt + "_edge_types.pt"),
                                         token=st.secrets["HF_TOKEN"])
     return embed_path, relation_weights_path, edge_types_path
 

requirements.txt

@@ -8,5 +8,4 @@ torch
 altair<5
 gspread
 oauth2client
-huggingface_hub
-matplotlib
+huggingface_hub