Merge branch 'main' into fix_llama_guard_inference

llama_toolchain/safety/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.

llama_toolchain/safety/api/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.

llama_toolchain/safety/api/config.py

@@ -1,3 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.
+
 from typing import List, Optional
 
 from pydantic import BaseModel

llama_toolchain/safety/api/datatypes.py

@@ -1,3 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.
+
 from enum import Enum
 from typing import Dict, Optional, Union
 

llama_toolchain/safety/shields/__init__.py

@@ -1,3 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.
+
 # supress warnings and spew of logs from hugging face
 import transformers
 

llama_toolchain/safety/shields/base.py

@@ -1,3 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.
+
 from abc import ABC, abstractmethod
 from typing import List, Union
 

llama_toolchain/safety/shields/code_scanner.py

@@ -1,3 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.
+
 from codeshield.cs import CodeShield
 from termcolor import cprint
 

llama_toolchain/safety/shields/contrib/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.

llama_toolchain/safety/shields/contrib/third_party_shield.py

@@ -1,3 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.
+
 import sys
 from typing import List
 
@@ -5,7 +17,11 @@ from llama_models.llama3_1.api.datatypes import Message
 
 parent_dir = "../.."
 sys.path.append(parent_dir)
-from llama_toolchain.safety.shields.base import OnViolationAction, ShieldBase, ShieldResponse
+from llama_toolchain.safety.shields.base import (
+    OnViolationAction,
+    ShieldBase,
+    ShieldResponse,
+)
 
 _INSTANCE = None
 

llama_toolchain/safety/shields/llama_guard.py

@@ -1,3 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.
+
 import re
 
 from string import Template
@@ -100,7 +112,7 @@ class LlamaGuardShield(ShieldBase):
     def instance(
         on_violation_action=OnViolationAction.RAISE,
         model_dir: str = None,
-        excluded_categories: List[str] = [],
+        excluded_categories: List[str] = None,
         disable_input_check: bool = False,
         disable_output_check: bool = False,
     ) -> "LlamaGuardShield":
@@ -119,7 +131,7 @@ class LlamaGuardShield(ShieldBase):
         self,
         on_violation_action: OnViolationAction = OnViolationAction.RAISE,
         model_dir: str = None,
-        excluded_categories: List[str] = [],
+        excluded_categories: List[str] = None,
         disable_input_check: bool = False,
         disable_output_check: bool = False,
     ):
@@ -129,6 +141,8 @@ class LlamaGuardShield(ShieldBase):
 
         assert model_dir is not None, "Llama Guard model_dir is None"
 
+        if excluded_categories is None:
+            excluded_categories = []
         assert len(excluded_categories) == 0 or all(
             x in SAFETY_CATEGORIES_TO_CODE_MAP.values() for x in excluded_categories
         ), "Invalid categories in excluded categories. Expected format is ['S1', 'S2', ..]"

llama_toolchain/safety/shields/prompt_guard.py

@@ -1,3 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.
+
 from enum import auto, Enum
 from typing import List
 

llama_toolchain/safety/shields/shield_runner.py

@@ -1,3 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described found in the
+# LICENSE file in the root directory of this source tree.
+
 import asyncio
 from typing import List
 
@@ -6,7 +18,7 @@ from llama_models.llama3_1.api.datatypes import Message, Role
 from .base import OnViolationAction, ShieldBase, ShieldResponse
 
 
-class SafetyException(Exception):
+class SafetyException(Exception):  # noqa: N818
     def __init__(self, response: ShieldResponse):
         self.response = response
         super().__init__(response.violation_return_message)