From e8be802310dd07b2e7b6216454f8273e6d470230 Mon Sep 17 00:00:00 2001 From: bkb2135 Date: Thu, 25 Jul 2024 19:01:19 +0000 Subject: [PATCH] Fix Dataset Creation --- prompting/conversation.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/prompting/conversation.py b/prompting/conversation.py index bd292ff4b..84a3bc359 100644 --- a/prompting/conversation.py +++ b/prompting/conversation.py @@ -41,8 +41,6 @@ def create_task( dataset = DATASETS.get(dataset_name, None) if dataset is None: raise ValueError(f"Dataset {dataset_name} not found") - elif task_name == SummarizationTask.name: - dataset = dataset(selector = "all") else: dataset = dataset() @@ -51,6 +49,12 @@ def create_task( translation_pipeline=translation_pipeline, context=dataset.next() ) + elif task_name == SummarizationTask.name: + return task( + llm_pipeline=llm_pipeline, + context=dataset.next(selector = "all"), + create_reference=create_reference, + ) return task( llm_pipeline=llm_pipeline,