code_instruct_example.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158
  1. # Copyright (c) Meta Platforms, Inc. and affiliates.
  2. # This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.
  3. # from accelerate import init_empty_weights, load_checkpoint_and_dispatch
  4. import fire
  5. import os
  6. import sys
  7. import time
  8. import torch
  9. from transformers import AutoTokenizer
  10. from llama_recipes.inference.safety_utils import get_safety_checker
  11. from llama_recipes.inference.model_utils import load_model, load_peft_model
  12. def handle_safety_check(are_safe_user_prompt, user_prompt, safety_results_user_prompt, are_safe_system_prompt, system_prompt, safety_results_system_prompt):
  13. """
  14. Handles the output based on the safety check of both user and system prompts.
  15. Parameters:
  16. - are_safe_user_prompt (bool): Indicates whether the user prompt is safe.
  17. - user_prompt (str): The user prompt that was checked for safety.
  18. - safety_results_user_prompt (list of tuples): A list of tuples for the user prompt containing the method, safety status, and safety report.
  19. - are_safe_system_prompt (bool): Indicates whether the system prompt is safe.
  20. - system_prompt (str): The system prompt that was checked for safety.
  21. - safety_results_system_prompt (list of tuples): A list of tuples for the system prompt containing the method, safety status, and safety report.
  22. """
  23. def print_safety_results(are_safe_prompt, prompt, safety_results, prompt_type="User"):
  24. """
  25. Prints the safety results for a prompt.
  26. Parameters:
  27. - are_safe_prompt (bool): Indicates whether the prompt is safe.
  28. - prompt (str): The prompt that was checked for safety.
  29. - safety_results (list of tuples): A list of tuples containing the method, safety status, and safety report.
  30. - prompt_type (str): The type of prompt (User/System).
  31. """
  32. if are_safe_prompt:
  33. print(f"{prompt_type} prompt deemed safe.")
  34. print(f"{prompt_type} prompt:\n{prompt}")
  35. else:
  36. print(f"{prompt_type} prompt deemed unsafe.")
  37. for method, is_safe, report in safety_results:
  38. if not is_safe:
  39. print(method)
  40. print(report)
  41. print(f"Skipping the inference as the {prompt_type.lower()} prompt is not safe.")
  42. sys.exit(1)
  43. # Check user prompt
  44. print_safety_results(are_safe_user_prompt, user_prompt, safety_results_user_prompt, "User")
  45. # Check system prompt
  46. print_safety_results(are_safe_system_prompt, system_prompt, safety_results_system_prompt, "System")
  47. def main(
  48. model_name,
  49. peft_model: str=None,
  50. quantization: bool=False,
  51. max_new_tokens =100, #The maximum numbers of tokens to generate
  52. prompt_file: str=None,
  53. seed: int=42, #seed value for reproducibility
  54. do_sample: bool=True, #Whether or not to use sampling ; use greedy decoding otherwise.
  55. min_length: int=None, #The minimum length of the sequence to be generated, input prompt + min_new_tokens
  56. use_cache: bool=False, #[optional] Whether or not the model should use the past last key/values attentions Whether or not the model should use the past last key/values attentions (if applicable to the model) to speed up decoding.
  57. top_p: float=0.9, # [optional] If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation.
  58. temperature: float=0.6, # [optional] The value used to modulate the next token probabilities.
  59. top_k: int=50, # [optional] The number of highest probability vocabulary tokens to keep for top-k-filtering.
  60. repetition_penalty: float=1.0, #The parameter for repetition penalty. 1.0 means no penalty.
  61. length_penalty: int=1, #[optional] Exponential penalty to the length that is used with beam-based generation.
  62. enable_azure_content_safety: bool=False, # Enable safety check with Azure content safety api
  63. enable_sensitive_topics: bool=False, # Enable check for sensitive topics using AuditNLG APIs
  64. enable_salesforce_content_safety: bool=True, # Enable safety check with Salesforce safety flan t5
  65. enable_llamaguard_content_safety: bool=False, # Enable safety check with Llama-Guard
  66. use_fast_kernels: bool = True, # Enable using SDPA from PyTroch Accelerated Transformers, make use Flash Attention and Xformer memory-efficient kernels
  67. **kwargs
  68. ):
  69. system_prompt = input("Please insert your system prompt: ")
  70. user_prompt = input("Please insert your prompt: ")
  71. chat = [
  72. {"role": "system", "content": system_prompt},
  73. {"role": "user", "content": user_prompt},
  74. ]
  75. # Set the seeds for reproducibility
  76. torch.cuda.manual_seed(seed)
  77. torch.manual_seed(seed)
  78. model = load_model(model_name, quantization)
  79. if peft_model:
  80. model = load_peft_model(model, peft_model)
  81. model.eval()
  82. if use_fast_kernels:
  83. """
  84. Setting 'use_fast_kernels' will enable
  85. using of Flash Attention or Xformer memory-efficient kernels
  86. based on the hardware being used. This would speed up inference when used for batched inputs.
  87. """
  88. try:
  89. from optimum.bettertransformer import BetterTransformer
  90. model = BetterTransformer.transform(model)
  91. except ImportError:
  92. print("Module 'optimum' not found. Please install 'optimum' it before proceeding.")
  93. tokenizer = AutoTokenizer.from_pretrained(model_name)
  94. safety_checker = get_safety_checker(enable_azure_content_safety,
  95. enable_sensitive_topics,
  96. enable_salesforce_content_safety,
  97. enable_llamaguard_content_safety,
  98. )
  99. # Safety check of the user prompt
  100. safety_results_user_prompt = [check(user_prompt) for check in safety_checker]
  101. safety_results_system_prompt = [check(system_prompt) for check in safety_checker]
  102. are_safe_user_prompt = all([r[1] for r in safety_results_user_prompt])
  103. are_safe_system_prompt = all([r[1] for r in safety_results_system_prompt])
  104. handle_safety_check(are_safe_user_prompt, user_prompt, safety_results_user_prompt, are_safe_system_prompt, system_prompt, safety_results_system_prompt)
  105. inputs = tokenizer.apply_chat_template(chat, return_tensors="pt").to("cuda")
  106. start = time.perf_counter()
  107. with torch.no_grad():
  108. outputs = model.generate(
  109. input_ids=inputs,
  110. max_new_tokens=max_new_tokens,
  111. do_sample=do_sample,
  112. top_p=top_p,
  113. temperature=temperature,
  114. min_length=min_length,
  115. use_cache=use_cache,
  116. top_k=top_k,
  117. repetition_penalty=repetition_penalty,
  118. length_penalty=length_penalty,
  119. **kwargs
  120. )
  121. e2e_inference_time = (time.perf_counter()-start)*1000
  122. print(f"the inference time is {e2e_inference_time} ms")
  123. output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
  124. # Safety check of the model output
  125. safety_results = [check(output_text) for check in safety_checker]
  126. are_safe = all([r[1] for r in safety_results])
  127. if are_safe:
  128. print("User input and model output deemed safe.")
  129. print(f"Model output:\n{output_text}")
  130. else:
  131. print("Model output deemed unsafe.")
  132. for method, is_safe, report in safety_results:
  133. if not is_safe:
  134. print(method)
  135. print(report)
  136. if __name__ == "__main__":
  137. fire.Fire(main)