Update instruction_template_retriever.py

instruction_template_retriever.py

@@ -1,5 +1,7 @@
 import itertools
 import json
+import pickle
+from random import Random
 
 from datasets import load_dataset
 import faiss
@@ -168,7 +170,9 @@ def unuse_gaussian_coverage_pooling(m):
 
 class InstructionTemplateRetriever:
     FINETEMPLATES_REVISION = "831ab22c90f9da011bd972585afdf609f40fa54b"
-    RETRIEVAL_EMBEDDING_NAME = "fineinstructions/instruction_template_retrieval_embedding"
+    RETRIEVAL_EMBEDDING_NAME = (
+        "fineinstructions/instruction_template_retrieval_embedding"
+    )
     RETRIEVAL_EMBEDDING_REVISION = "db4efbde126216250ffa5a356663fc7da3bf7856"
 
     def __init__(
@@ -222,6 +226,21 @@ class InstructionTemplateRetriever:
         elif torch.backends.mps.is_available():
             self.m = self.m.to("mps")
 
+        with open(
+            hf_hub_download(
+                "fineinstructions/finetemplates",
+                "faiss_index/reweighting_stats.pkl",
+                revision=FINETEMPLATES_REVISION,
+                repo_type="dataset",
+            ),
+            "rb",
+        ) as reweighting_stats_fp:
+            reweighting_stats = pickle.load(reweighting_stats_fp)
+        self.resampling_weights = reweighting_stats["resampling_weights"]
+        self.template_variable_count_mapping = reweighting_stats[
+            "template_variable_count_mapping"
+        ]
+
     def _filter_rows(self, rows, filter_string):
         if not rows:
             return []
@@ -233,7 +252,14 @@ class InstructionTemplateRetriever:
             return rows
 
     def search(
-        self, document, filters="", search_k=20000, max_results=250, deduplicate=True
+        self,
+        document,
+        filters="",
+        search_k=20000,
+        max_results=250,
+        deduplicate=True,
+        reweight=False,
+        reweighting_epsilon=True,
     ):
         """
         Given a document
@@ -246,6 +272,31 @@ class InstructionTemplateRetriever:
             deduplicate (bool): Deduplicate results between coverage sections.
         """
 
+        def _reweight(inp, k=None):
+            if reweight:
+                inp0, inp = itertools.tee(inp)
+                first_row = next(inp0)
+                r = Random(first_row[1].item())
+                epsilon = reweighting_epsilon
+                bucket = first_row[1]
+                items = []
+                weights = []
+                for i, s in inp:
+                    if abs(bucket - s.item()) <= epsilon:
+                        items.append((i, s))
+                        weights.append(
+                            self.resampling_weights[
+                                self.template_variable_count_mapping[i.item()]
+                            ]
+                        )
+                    else:
+                        break
+                return r.choices(
+                    items, weights=weights, k=(len(items) if k is None else k)
+                )
+            else:
+                return inp
+
         # Search FAISS index
         vecs = self.m.encode([document], normalize_embeddings=False).reshape(
             -1, self.m[0].auto_model.config.hidden_size
@@ -284,7 +335,7 @@ class InstructionTemplateRetriever:
                     "score": s.item(),
                     **d_in_mem[i.item()],
                 }
-                for i, s in zip(indices, scores)
+                for i, s in _reweight(zip(indices, scores), k=None)
             ]
             for chunk_idx, (indices, scores) in enumerate(
                 zip(indices_per_input, scores_per_input)