ensure ksims and kscores are correct size

app.py

@@ -57,11 +57,12 @@ def knn_get_score(knn, k, cat, vec):
     sortinds = np.flip(np.argsort(cos_sim_table, axis=1), axis=1)
     if debug: st.write('sortinds.shape', sortinds.shape)
     # Get corresponding scores for the sorted vectors
-    kscores = scores[sortinds][:k]
+    kscores = scores[sortinds][:,:k]
     if debug: st.write('kscores.shape', kscores.shape)
     # Get actual sorted similiarity scores
     # (line copied from clip_retrieval_knn.py even though sortinds.shape[0] == 1 here)
     ksims = cos_sim_table[np.expand_dims(np.arange(sortinds.shape[0]), axis=1), sortinds]
+    ksims = ksims[:,:k]
     if debug: st.write('ksims.shape', ksims.shape)
     # Apply normalization after exponential formula
     ksims = softmax(10**ksims)