fix: optimise softmax calculation in CustomLID

[kshkrvea]

Optimized prediction inference to boost overall throughput.

This refactor gets a slight speedup while keeping the outputs exactly the same as the original method.

What I did:

1. Sorted only top-k part of the raw logits -> reduce sorting step complexity from O(n*log(n)) to O(n) where n is a number of classes.
2. Restricted softmax calculations exclusively to top-k classes. Denominator kept same.

Numerical outputs of the functions are same, while computation time (averaged over 10 runs) dropped down by 10-ish %. You can see it in the first commit c14a610 of this PR in report.txt or bellow:

[before vs optimized_before]
Max Absolute Difference: 0.00e+00
Speed Original: 2508.8 samples/sec, 0.000399 sec/sample
Speed Optimized: 2879.5 samples/sec, 0.000347 sec/sample
Total Speedup: 12.9%

[after vs optimized_after]
Max Absolute Difference: 0.00e+00
Speed Original: 3627.7 samples/sec, 0.000276 sec/sample
Speed Optimized: 4463.0 samples/sec, 0.000224 sec/sample
Total Speedup: 18.7%

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

This PR refactors top‑k selection in LID inference to avoid full sorting and to compute probabilities for only the requested top‑k outputs.

Changes:

• Uses np.argpartition + partial sort to compute top‑k indices more efficiently.
• Refactors softmax math in predict_limit_before_softmax to compute top‑k probabilities without materializing the full softmax vector.
• Applies the same top‑k selection approach in predict_limit_after_softmax.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.