From c9d9035f9507588d3c4b9d21ca93648332244f6c Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <alvaro@argilla.io>
Date: Tue, 21 Nov 2023 11:47:21 +0100
Subject: [PATCH] Fix `apply_chat_template` function for `dpo` and unknown
 `task` (#30)

* Fix `apply_chat_template` function for `dpo` and unknown `task`

* Remove duplicated `# coding=utf-8`

* Manually run `black --line-length 119`
---
 src/alignment/configs.py     |  1 -
 src/alignment/data.py        | 12 +++++++-----
 src/alignment/model_utils.py |  1 -
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/alignment/configs.py b/src/alignment/configs.py
index d785e16..9097d94 100644
--- a/src/alignment/configs.py
+++ b/src/alignment/configs.py
@@ -1,5 +1,4 @@
 # coding=utf-8
-# coding=utf-8
 # Copyright 2023 The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
diff --git a/src/alignment/data.py b/src/alignment/data.py
index 838169a..da4b979 100644
--- a/src/alignment/data.py
+++ b/src/alignment/data.py
@@ -1,5 +1,4 @@
 # coding=utf-8
-# coding=utf-8
 # Copyright 2023 The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -72,12 +71,15 @@ def apply_chat_template(
             example["text_prompt"] = tokenizer.apply_chat_template(
                 prompt_messages, tokenize=False, add_generation_prompt=True
             )
-
-        example["text_chosen"] = _strip_prefix(example["text_chosen"], assistant_prefix)
-        example["text_rejected"] = _strip_prefix(example["text_rejected"], assistant_prefix)
+            example["text_chosen"] = _strip_prefix(example["text_chosen"], assistant_prefix)
+            example["text_rejected"] = _strip_prefix(example["text_rejected"], assistant_prefix)
+        else:
+            raise ValueError(
+                f"Could not format example as dialogue for `dpo` task! Require `[chosen, rejected]` keys but found {list(example.keys())}"
+            )
     else:
         raise ValueError(
-            f"Could not format example as dialogue for `dpo` task! Require `[chosen, rejected]` keys but found {list(example.keys())}"
+            f"Task {task} not supported, please ensure that the provided task is one of {['sft', 'generation', 'rm', 'dpo']}"
         )
     return example
 
diff --git a/src/alignment/model_utils.py b/src/alignment/model_utils.py
index 9463f2e..b9d2315 100644
--- a/src/alignment/model_utils.py
+++ b/src/alignment/model_utils.py
@@ -1,5 +1,4 @@
 # coding=utf-8
-# coding=utf-8
 # Copyright 2023 The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");