Merge pull request #28 from opentensor/topk

ensure topk isn't out of bounds
opentensor · Oct 2, 2023 · d6e2919 · d6e2919
2 parents 6322c96 + e2e348f
commit d6e2919
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/prompting/validators/reward/diversity.py b/prompting/validators/reward/diversity.py
@@ -129,9 +129,10 @@ def regularise(rewards):
         )
 
         # Reward to be at the bottom_k smallest of the 1 - similarity score.
-        rewards = torch.topk(
-            (1 - torch.abs(similarity)), self.history_reward_bottom_k, largest=False
-        )[0][:, -1]
+        bottom_k = min(self.history_reward_bottom_k, len(similarity))
+        rewards = torch.topk((1 - torch.abs(similarity)), bottom_k, largest=False)[0][
+            :, -1
+        ]
 
         return regularise(rewards)
 
@@ -144,9 +145,10 @@ def regularise(rewards):
         similarity = pairwise_cosine_similarity(embeddings, embeddings)
 
         # Reward to be at the 10% quantile of the 1 - similarity score.
-        rewards = torch.topk(
-            (1 - torch.abs(similarity)), self.reward_bottom_k, largest=False
-        )[0][:, -1]
+        bottom_k = min(self.reward_bottom_k, len(similarity))
+        rewards = torch.topk((1 - torch.abs(similarity)), bottom_k, largest=False)[0][
+            :, -1
+        ]
 
         return regularise(rewards)