mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-28 23:07:11 +00:00
db9d5b213a
Optimizing important numerical code and making it run faster.
Performance went up by 1.48x (148%). Runtime went down from 138715us to
56020us
Optimization explanation:
The `cosine_similarity_top_k` function is where we made the most
significant optimizations.
Instead of sorting the entire score_array which needs considering all
elements, `np.argpartition` is utilized to find the top_k largest scores
indices, this operation has a time complexity of O(n), higher
performance than sorting. Remember, `np.argpartition` doesn't guarantee
the order of the values. So we need to use argsort() to get the indices
that would sort our top-k values after partitioning, which is much more
efficient because it only sorts the top-K elements, not the entire
array. Then to get the row and column indices of sorted top_k scores in
the original score array, we use `np.unravel_index`. This operation is
more efficient and cleaner than a list comprehension.
The code has been tested for correctness by running the following
snippet on both the original function and the optimized function and
averaged over 5 times.
```
def test_cosine_similarity_top_k_large_matrices():
X = np.random.rand(1000, 1000)
Y = np.random.rand(1000, 1000)
top_k = 100
score_threshold = 0.5
gc.disable()
counter = time.perf_counter_ns()
return_value = cosine_similarity_top_k(X, Y, top_k, score_threshold)
duration = time.perf_counter_ns() - counter
gc.enable()
```
@hwaking @hwchase17 @jerwelborn
Unit tests pass, I also generated more regression tests which all
passed.