Refactor get_question method in Context class

EvolvingLMMs-Lab · Mar 30, 2024 · 8d56856 · 8d56856
1 parent 9979b03
commit 8d56856
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lmms_eval/api/samplers.py b/lmms_eval/api/samplers.py
@@ -27,7 +27,8 @@ def __init__(self, task, few_shot_delimiter: str = "\n\n", target_delimiter: str
         self.contexts = []
 
     def get_question(self, doc, model_specific_prompt_kwargs=None):
-        return self.doc_to_text(doc, model_specific_prompt_kwargs) if (self.doc_to_choice is None or type(self.doc_to_text(doc)) is str) else self.doc_to_choice(doc)[self.doc_to_text(doc)]
+        text = self.doc_to_text(doc, model_specific_prompt_kwargs)
+        return text if (self.doc_to_choice is None or isinstance(text, str)) else self.doc_to_choice(doc)[text]
 
     def get_target(self, doc):
         return (