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