SN1-212: Moving from github to huggingface-github dataset

macrocosm-os · Sep 3, 2024 · 002e2af · 002e2af
1 parent d816bc3
commit 002e2af
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/prompting/tasks/programming_task.py b/prompting/tasks/programming_task.py
@@ -1,5 +1,9 @@
 from typing import ClassVar
-from prompting.datasets.github import GithubDatasetEntry, MIN_INPUT_LINES, OUTPUT_LINES
+from prompting.datasets.huggingface_github import (
+    HuggingFaceGithubDatasetEntry,
+    MIN_INPUT_LINES,
+    OUTPUT_LINES,
+)
 from prompting.tasks.base_task import BaseTextTask
 from prompting.rewards.reward import BaseRewardConfig, WeightedRewardModel
 from prompting.rewards.rouge import RougeRewardModel
@@ -36,7 +40,7 @@ class ProgrammingTask(BaseTextTask):
     query: str | None = None
     reference: str | None = None
 
-    def make_query(self, dataset_entry: GithubDatasetEntry):
+    def make_query(self, dataset_entry: HuggingFaceGithubDatasetEntry):
         modified_code = model_manager.generate(
             [CODE_MODIFICATION_PROMPT.format(file_content=dataset_entry.file_content)]
         )[0]