Update controllable_blender/generation_methods.py

controllable_blender/generation_methods.py

@@ -18,7 +18,7 @@ class VocabTopKSampling(TopKSampling):
                  k: int,
                  wordlist: Wordlist,
                  *args, **kwargs):
-        super().__init__(*args, **kwargs)
+        super().__init__(k=k, *args, **kwargs)
         self.k = k
         self.wordlist = wordlist
 
@@ -44,7 +44,7 @@ class RerankedTopKSampling(TreeSearch):
                  reranker: Reranker,
                  tokenids_to_text: Callable,
                  *args, **kwargs):
-        super().__init__(*args, **kwargs)
+        super().__init__(k=k, *args, **kwargs)
         self.k = k
         self.reranker = reranker
         self.tokenids_to_text = tokenids_to_text