Fix lint issues

Stability-AI · Oct 20, 2023 · b46ea45 · b46ea45
1 parent 3820f10
commit b46ea45
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 11 deletions.
diff --git a/lm_eval/base.py b/lm_eval/base.py
@@ -519,8 +519,7 @@ def _process_doc(self, doc):
         return doc
 
     def fewshot_examples(self, k, rnd, stratified=False):
-        """Returns few shot examples from training docs
-        """
+        """Returns few shot examples from training docs"""
         if self._training_docs is None:
             self._training_docs = list(self.training_docs())
 
@@ -559,8 +558,7 @@ def _stratified_fewshot_examples(self, docs, k, rnd):
         # but this may not be guaranteed, so calculate the number of sample
         # for each target per method call
         target_to_num_samples = {
-            target: int(ratio * k)
-            for target, ratio in self._target_to_ratio.items()
+            target: int(ratio * k) for target, ratio in self._target_to_ratio.items()
         }
         # Handle any rounding discrepancies by adjusting the counts
         remaining_samples = k - sum(target_to_num_samples.values())
@@ -649,7 +647,13 @@ def fewshot_description(self):
 
     @utils.positional_deprecated
     def fewshot_context(
-        self, doc, num_fewshot, provide_description=None, rnd=None, description=None, stratified=False
+        self,
+        doc,
+        num_fewshot,
+        provide_description=None,
+        rnd=None,
+        description=None,
+        stratified=False,
     ):
         """Returns a fewshot context string that is made up of a prepended description
         (if provided), the `num_fewshot` number of examples, and an appended prompt example.
@@ -702,7 +706,9 @@ def fewshot_context(
         else:
             # for sets with no training docs, draw from other set *but ensure no overlap with current doc*
             if self.has_training_docs():
-                fewshotex = self.fewshot_examples(k=num_fewshot, rnd=rnd, stratified=stratified)
+                fewshotex = self.fewshot_examples(
+                    k=num_fewshot, rnd=rnd, stratified=stratified
+                )
             else:
                 if self._fewshot_docs is None:
                     self._fewshot_docs = list(
@@ -712,7 +718,9 @@ def fewshot_context(
                     )
 
                 if stratified:
-                    fewshotex = self._stratified_fewshot_examples(self._fewshot_docs, num_fewshot + 1, rnd=rnd)
+                    fewshotex = self._stratified_fewshot_examples(
+                        self._fewshot_docs, num_fewshot + 1, rnd=rnd
+                    )
                 else:
                     fewshotex = rnd.sample(self._fewshot_docs, num_fewshot + 1)
 

diff --git a/lm_eval/tasks/ja/jcola.py b/lm_eval/tasks/ja/jcola.py
@@ -43,9 +43,19 @@ def construct_requests(self, doc, ctx):
         ll_false, _ = rf.loglikelihood(ctx, " %s" % self.CHOICES[0])
         return ll_true, ll_false
 
-    def fewshot_context(self, doc, num_fewshot, provide_description=None, rnd=None, description=None, stratified=False):
+    def fewshot_context(
+        self,
+        doc,
+        num_fewshot,
+        provide_description=None,
+        rnd=None,
+        description=None,
+        stratified=False,
+    ):
         # Use stratified sampling
-        return super().fewshot_context(doc, num_fewshot, provide_description, rnd, description, stratified=True)
+        return super().fewshot_context(
+            doc, num_fewshot, provide_description, rnd, description, stratified=True
+        )
 
 
 class JCoLAWithJAAlpacaPrompt(JCoLA):

diff --git a/lm_eval/tasks/ja/jnli.py b/lm_eval/tasks/ja/jnli.py
@@ -92,7 +92,15 @@ def construct_requests(self, doc, ctx):
             lls.append(rf.greedy_until(ctx, [self.SEP]))
         return lls
 
-    def fewshot_context(self, doc, num_fewshot, provide_description=None, rnd=None, description=None, stratified=False):
+    def fewshot_context(
+        self,
+        doc,
+        num_fewshot,
+        provide_description=None,
+        rnd=None,
+        description=None,
+        stratified=False,
+    ):
         """
         TODO: move this to `MultipleChoiceTask`.
         Directly implementing this in `MultipleChoiceTask` will break the task versioning
@@ -101,7 +109,9 @@ def fewshot_context(self, doc, num_fewshot, provide_description=None, rnd=None,
         only after all tasks have been updated, then we can move this to `MultipleChoiceTask`.
         """
         # Use stratified sampling
-        return super().fewshot_context(doc, num_fewshot, provide_description, rnd, description, stratified=True)
+        return super().fewshot_context(
+            doc, num_fewshot, provide_description, rnd, description, stratified=True
+        )
 
 
 class JNLIWithJAAlpacaPrompt(JNLIWithFintanPrompt):