Consider batch_size when sorting similar vectors (#4388)

This commit is contained in:
adrianeboyd 2019-10-07 13:38:35 +02:00 committed by Matthew Honnibal
parent a3509f67d4
commit d53a8d9313

View File

@ -337,7 +337,7 @@ cdef class Vectors:
scores[i:i+batch_size] = xp.partition(sims, -n, axis=1)[:,-n:] scores[i:i+batch_size] = xp.partition(sims, -n, axis=1)[:,-n:]
if sort: if sort:
sorted_index = xp.arange(scores.shape[0])[:,None],xp.argsort(scores, axis=1)[:,::-1] sorted_index = xp.arange(scores.shape[0])[:,None],xp.argsort(scores[i:i+batch_size], axis=1)[:,::-1]
scores[i:i+batch_size] = scores[sorted_index] scores[i:i+batch_size] = scores[sorted_index]
best_rows[i:i+batch_size] = best_rows[sorted_index] best_rows[i:i+batch_size] = best_rows[sorted_index]