Removing Wiki & MultiChoice, rebalancing probabilities

prompting/tasks/multi_step_reasoning.py

@@ -6,7 +6,7 @@
 from prompting.datasets.random_website import DDGDatasetEntry
 from prompting.rewards.relevance import RelevanceRewardModel
 from prompting.rewards.reward import BaseRewardConfig, BaseRewardModel
-from prompting.tasks.qa import WikiQuestionAnsweringTask
+from prompting.tasks.qa import WebQuestionAnsweringTask
 from shared.base import Context
 from validator_api.test_time_inference import generate_response
 
@@ -71,7 +71,7 @@ class MultiStepReasoningRewardConfig(BaseRewardConfig):
 ]
 
 
-class MultiStepReasoningTask(WikiQuestionAnsweringTask):
+class MultiStepReasoningTask(WebQuestionAnsweringTask):
     """QuestionAnsweringTasks must be initialised with an LLM pipeline to generate query and reference plus
     context from a dataset to base the query on"""
 

prompting/tasks/qa.py

@@ -46,29 +46,6 @@ class QARewardConfig(BaseRewardConfig):
     ]
     penalty_definition: ClassVar[list[BaseRewardModel]] = [RougeRewardModel(weight=0.5)]
 
-
-class WikiQuestionAnsweringTask(BaseTextTask):
-    """QuestionAnsweringTasks must be initialised with an LLM pipeline to generate query and reference plus
-    context from a dataset to base the query on"""
-
-    name: ClassVar[str] = "wiki_qa"
-    query_system_prompt: ClassVar[str] = QUERY_SYSTEM_PROMPT
-    reference_system_prompt: ClassVar[str] = REFERENCE_SYSTEM_PROMPT
-    augmentation_system_prompt: ClassVar[str] = ""
-    query: str | None = None
-    reference: str | None = None
-
-    def make_query(self, dataset_entry: Context):
-        query_prompt = QUERY_PROMPT_TEMPLATE.format(context=dataset_entry.content)
-        self.query = self.generate_query(messages=[query_prompt])
-        return self.query
-
-    async def make_reference(self, dataset_entry: Context):
-        reference_prompt = REFERENCE_PROMPT_TEMPLATE.format(context=dataset_entry.content, question=self.query)
-        self.reference = self.generate_reference(messages=[{"role": "user", "content": reference_prompt}])
-        return self.reference
-
-
 class WebQuestionAnsweringTask(BaseTextTask):
     """QuestionAnsweringTasks must be initialised with an LLM pipeline to generate query and reference plus
     context from a dataset to base the query on"""

prompting/tasks/task_registry.py

@@ -12,10 +12,9 @@
 from prompting.rewards.reward import BaseRewardConfig
 from prompting.tasks.base_task import BaseTextTask
 from prompting.tasks.inference import InferenceRewardConfig, InferenceTask
-from prompting.tasks.multi_choice import MultiChoiceRewardConfig, MultiChoiceTask
 from prompting.tasks.multi_step_reasoning import MultiStepReasoningRewardConfig, MultiStepReasoningTask
 from prompting.tasks.programming_task import ProgrammingRewardConfig, ProgrammingTask
-from prompting.tasks.qa import QARewardConfig, WebQuestionAnsweringTask, WikiQuestionAnsweringTask
+from prompting.tasks.qa import QARewardConfig, WebQuestionAnsweringTask
 from prompting.tasks.web_retrieval import WebRetrievalRewardConfig, WebRetrievalTask
 from shared.base import BaseDataset
 
@@ -34,22 +33,13 @@ def __hash__(self):
 
 class TaskRegistry(BaseModel):
     task_configs: ClassVar[list[TaskConfig]] = [
-        TaskConfig(
-            task=WikiQuestionAnsweringTask, probability=0.05, datasets=[WikiDataset], reward_model=QARewardConfig
-        ),
-        TaskConfig(task=WebQuestionAnsweringTask, probability=0.15, datasets=[DDGDataset], reward_model=QARewardConfig),
+        TaskConfig(task=WebQuestionAnsweringTask, probability=0.05, datasets=[DDGDataset], reward_model=QARewardConfig),
         TaskConfig(
             task=InferenceTask,
             probability=0.3,
             datasets=[SN13Dataset],
             reward_model=InferenceRewardConfig,
         ),
-        TaskConfig(
-            task=MultiChoiceTask,
-            probability=0.2,
-            datasets=[WikiDataset],
-            reward_model=MultiChoiceRewardConfig,
-        ),
         TaskConfig(
             task=ProgrammingTask,
             probability=0.10,
@@ -58,13 +48,13 @@ class TaskRegistry(BaseModel):
         ),
         TaskConfig(
             task=WebRetrievalTask,
-            probability=0.1,
+            probability=0.25,
             datasets=[DDGDataset],
             reward_model=WebRetrievalRewardConfig,
         ),
         TaskConfig(
             task=MultiStepReasoningTask,
-            probability=0.1,
+            probability=0.3,
             datasets=[DDGDataset],
             reward_model=MultiStepReasoningRewardConfig,
         ),