code cleaning

dataloader/__init__.py

@@ -0,0 +1,10 @@
+from agentverse.registry import Registry
+
+dataloader_registry = Registry(name="dataloader")
+
+from .gsm8k import GSM8KLoader
+from .responsegen import ResponseGenLoader
+from .humaneval import HumanevalLoader
+from .commongen import CommongenLoader
+from .mgsm import MGSMLoader
+from .logic_grid import LogicGridLoader

dataloader/commongen.py

@@ -0,0 +1,21 @@
+from .dataloader import DataLoader
+from . import dataloader_registry
+import json
+
+
+@dataloader_registry.register("tasksolving/commongen/gpt-4")
+@dataloader_registry.register("tasksolving/commongen/gpt-3.5")
+class CommongenLoader(DataLoader):
+    def __init__(self, path: str):
+        super().__init__(path)
+
+    def load(self):
+        with open(self.path) as f:
+            for line in f:
+                line = json.loads(line)
+                self.examples.append(
+                    {
+                        "input": line["concepts"],
+                        "answer": None,
+                    }
+                )

dataloader/dataloader.py

@@ -0,0 +1,19 @@
+import json
+from abc import abstractmethod
+
+
+class DataLoader:
+    def __init__(self, path: str):
+        self.path = path
+        self.examples = []
+        self.load()
+
+    @abstractmethod
+    def load(self):
+        """Make sure that each example is formatted as {"input": ..., "answer": ...}"""
+        with open(self.path) as f:
+            for line in f:
+                self.examples.append(json.loads(line))
+
+    def __iter__(self):
+        return iter(self.examples)

dataloader/gsm8k.py

@@ -0,0 +1,22 @@
+from .dataloader import DataLoader
+from . import dataloader_registry
+import json
+import re
+
+
+@dataloader_registry.register("tasksolving/gsm8k")
+class GSM8KLoader(DataLoader):
+    def __init__(self, path: str):
+        self.answer_pat = re.compile(r"#### (-?\d+)")
+        super().__init__(path)
+
+    def load(self):
+        with open(self.path) as f:
+            for line in f:
+                line = json.loads(line)
+                self.examples.append(
+                    {
+                        "input": line["question"],
+                        "answer": line["answer"].split('#### ')[-1],
+                    }
+                )

dataloader/humaneval.py

@@ -0,0 +1,21 @@
+from .dataloader import DataLoader
+from . import dataloader_registry
+import json
+
+
+@dataloader_registry.register("tasksolving/humaneval/gpt-4")
+@dataloader_registry.register("tasksolving/humaneval/gpt-3.5")
+class HumanevalLoader(DataLoader):
+    def __init__(self, path: str):
+        super().__init__(path)
+
+    def load(self):
+        with open(self.path) as f:
+            for line in f:
+                line = json.loads(line)
+                self.examples.append(
+                    {
+                        "input": line["prompt"],
+                        "answer": line["test"],
+                    }
+                )

dataloader/logic_grid.py

@@ -0,0 +1,22 @@
+from .dataloader import DataLoader
+from . import dataloader_registry
+import json
+import re
+
+
+@dataloader_registry.register("tasksolving/logic_grid/gpt-4")
+class LogicGridLoader(DataLoader):
+    def __init__(self, path: str):
+        self.answer_pat = re.compile(r"#### (-?\d+)")
+        super().__init__(path)
+
+    def load(self):
+        with open(self.path) as f:
+            for line in f:
+                line = json.loads(line)
+                self.examples.append(
+                    {
+                        "input": line["inputs"],
+                        "answer": line["targets"][0],
+                    }
+                )

dataloader/mgsm.py

@@ -0,0 +1,23 @@
+from .dataloader import DataLoader
+from . import dataloader_registry
+import json
+import re
+
+
+@dataloader_registry.register("tasksolving/mgsm/gpt-4")
+@dataloader_registry.register("tasksolving/mgsm/gpt-3.5")
+class MGSMLoader(DataLoader):
+    def __init__(self, path: str):
+        self.answer_pat = re.compile(r"#### (-?\d+)")
+        super().__init__(path)
+
+    def load(self):
+        with open(self.path) as f:
+            for line in f:
+                line = json.loads(line)
+                self.examples.append(
+                    {
+                        "input": line["question"],
+                        "answer": line["answer_number"],
+                    }
+                )

dataloader/responsegen.py

@@ -0,0 +1,21 @@
+from .dataloader import DataLoader
+from . import dataloader_registry
+import json
+
+
+@dataloader_registry.register("tasksolving/responsegen/gpt-3.5")
+@dataloader_registry.register("tasksolving/responsegen/gpt-4")
+class ResponseGenLoader(DataLoader):
+    def __init__(self, path: str):
+        super().__init__(path)
+
+    def load(self):
+        with open(self.path) as f:
+            for line in f:
+                line = json.loads(line)
+                self.examples.append(
+                    {
+                        "input": line["input"],
+                        "answer": line["answer"],
+                    }
+                )