소스 검색

Add option to enable Llamaguard content safety check in chat_completion (#354)

Hamid Shojanazeri 1 년 전
부모
커밋
98fcc538ff
1개의 변경된 파일2개의 추가작업 그리고 0개의 파일을 삭제
  1. 2 0
      examples/chat_completion/chat_completion.py

+ 2 - 0
examples/chat_completion/chat_completion.py

@@ -35,6 +35,7 @@ def main(
     enable_sensitive_topics: bool=False, # Enable check for sensitive topics using AuditNLG APIs
     enable_saleforce_content_safety: bool=True, # Enable safety check woth Saleforce safety flan t5
     use_fast_kernels: bool = False, # Enable using SDPA from PyTorch Accelerated Transformers, make use Flash Attention and Xformer memory-efficient kernels
+    enable_llamaguard_content_safety: bool = False,
     **kwargs
 ):
     if prompt_file is not None:
@@ -90,6 +91,7 @@ def main(
             safety_checker = get_safety_checker(enable_azure_content_safety,
                                         enable_sensitive_topics,
                                         enable_saleforce_content_safety,
+                                        enable_llamaguard_content_safety,
                                         )
             # Safety check of the user prompt
             safety_results = [check(dialogs[idx][0]["content"]) for check in safety_checker]