Add klue-sts task

lm_eval/tasks/__init__.py

@@ -51,6 +51,7 @@ from . import blimp
 from . import asdiv
 from . import gsm8k
 from . import storycloze
+from . import klue
 
 ########################################
 # Translation tasks
@@ -302,6 +303,9 @@ TASK_REGISTRY = {
     # "storycloze_2016": storycloze.StoryCloze2016,
     # "storycloze_2018": storycloze.StoryCloze2018,
     # "sat": sat.SATAnalogies,
+    
+    # KLUE
+    "klue_sts": klue.STS
 }
 
 

lm_eval/tasks/klue.py

+"""
+NSMC:
+"""
+import numpy as np
+from lm_eval.base import rf, Task
+from lm_eval.metrics import mean, matthews_corrcoef, f1_score, yesno
+from lm_eval.utils import general_detokenize
+
+_CITATION = """
+@misc{park2021klue,
+      title={KLUE: Korean Language Understanding Evaluation},
+      author={Sungjoon Park and Jihyung Moon and Sungdong Kim and Won Ik Cho and Jiyoon Han and Jangwon Park and Chisung Song and Junseong Kim and Yongsook Song and Taehwan Oh and Joohong Lee and Juhyun Oh and Sungwon Lyu and Younghoon Jeong and Inkwon Lee and Sangwoo Seo and Dongjun Lee and Hyunwoo Kim and Myeonghwa Lee and Seongbo Jang and Seungwon Do and Sunkyoung Kim and Kyungtae Lim and Jongwon Lee and Kyumin Park and Jamin Shin and Seonghyun Kim and Lucy Park and Alice Oh and Jungwoo Ha and Kyunghyun Cho},
+      year={2021},
+      eprint={2105.09680},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+"""
+
+
+class STS(Task):
+    VERSION = 0
+    DATASET_PATH = "klue"
+    DATASET_NAME = "sts"
+
+    def has_training_docs(self):
+        return True
+
+    def has_validation_docs(self):
+        return True
+
+    def has_test_docs(self):
+        return False
+
+    def training_docs(self):
+        if self._training_docs is None:
+            self._training_docs = list(self.dataset["train"])
+        return self._training_docs
+
+    def validation_docs(self):
+        return self.dataset["validation"]
+
+    def doc_to_text(self, doc):
+        return "질문: 문장 1과 문장 2는 서로 유사한 의미를 가지나요?\n문장 1:{}\n문장 2:{}\n정답:".format(
+            general_detokenize(doc["sentence1"]),
+            general_detokenize(doc["sentence2"]) 
+        )
+
+    def doc_to_target(self, doc):
+        return " {}".format({1: " 예", 0: " 아니"}[doc["labels"]["binary-label"]])
+
+    def construct_requests(self, doc, ctx):
+        ll_positive, _ = rf.loglikelihood(ctx, " 예")
+        ll_negative, _ = rf.loglikelihood(ctx, " 아니")
+        return ll_positive, ll_negative
+
+    def process_results(self, doc, results):
+        ll_positive, ll_negative = results
+        pred = ll_positive > ll_negative
+        gold = doc["labels"]["binary-label"]
+        return {
+            "acc": pred == gold,
+            "f1": (gold, pred)
+        }
+
+    def higher_is_better(self):
+        return {
+            "acc": True,
+            "f1": True
+        }
+
+    def aggregation(self):
+        return {
+            "acc": mean,
+            "f1": f1_score
+        }
+        
\ No newline at end of file