|
@@ -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()
|