Browse Source

Fix merge conflict with main

Michael Tontchev 1 year ago
parent
commit
b14cb02272
1 changed files with 6 additions and 2 deletions
  1. 6 2
      tests/test_finetuning_data_formatter.py

+ 6 - 2
tests/test_finetuning_data_formatter.py

@@ -1,11 +1,10 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # This software may be used and distributed according to the terms of the Llama Guard Community License Agreement.
 
+from enum import Enum
 import unittest
 from typing import Optional
 
-from examples.llama_guard.prompt_format import AgentType
-
 from llama_recipes.data.llama_guard.finetuning_data_formatter import (
     AugmentationConfigs,
     Category,
@@ -19,6 +18,11 @@ from llama_recipes.data.llama_guard.finetuning_data_formatter import (
 )
 
 
+class AgentType(Enum):
+    AGENT = "Agent"
+    USER = "User"
+
+
 class FinetuningDataFormatterTests(unittest.TestCase):
     def setUp(self) -> None:
         super().setUp()