add size of databases

app.py

@@ -126,16 +126,19 @@ def retrieval():
         selected_database = st.selectbox(
                 'Select database',('Lenselink', 'Davis', 'DUD-E')
             )
+        l = {
+            'Lenselink': 314707,
+            'Davis': 30056,
+            'DUDE': 1434019,
+        }
         if selected_database == 'DUD-E':
             selected_database = 'DUDE'
-        c1, c2 = st.columns(2)
-        with c2:
-            st.image('figures/multi_drugs.png', use_column_width='always') #, width=125)
-        with st.spinner('Loading data...'):
+        st.image('figures/multi_drugs.png', use_column_width='always') #, width=125)
+        with st.spinner(f'Loading {l[selected_database]} drugs...'):
             batch_size = 2048
             dataset = DrugRetrieval(os.path.join(data_path, selected_database), sequence, query_embedding)
             dataloader = DataLoader(dataset, num_workers=2, batch_size=batch_size, shuffle=False, collate_fn=collate_target)
-        st.success('Data loaded.')
+        st.success(f'{l[selected_database]} drugs loaded.')
     
     with col4:
         selected_encoder = st.selectbox(