code_infilling_example.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  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 torch
  6. import os
  7. import sys
  8. import time
  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 main(
  13. model_name,
  14. peft_model: str=None,
  15. quantization: bool=False,
  16. max_new_tokens =100, #The maximum numbers of tokens to generate
  17. prompt_file: str=None,
  18. seed: int=42, #seed value for reproducibility
  19. do_sample: bool=True, #Whether or not to use sampling ; use greedy decoding otherwise.
  20. min_length: int=None, #The minimum length of the sequence to be generated, input prompt + min_new_tokens
  21. use_cache: bool=True, #[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.
  22. 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.
  23. temperature: float=0.6, # [optional] The value used to modulate the next token probabilities.
  24. top_k: int=50, # [optional] The number of highest probability vocabulary tokens to keep for top-k-filtering.
  25. repetition_penalty: float=1.0, #The parameter for repetition penalty. 1.0 means no penalty.
  26. length_penalty: int=1, #[optional] Exponential penalty to the length that is used with beam-based generation.
  27. enable_azure_content_safety: bool=False, # Enable safety check with Azure content safety api
  28. enable_sensitive_topics: bool=False, # Enable check for sensitive topics using AuditNLG APIs
  29. enable_salesforce_content_safety: bool=True, # Enable safety check with Salesforce safety flan t5
  30. enable_llamaguard_content_safety: bool=False, # Enable safety check with Llama-Guard
  31. use_fast_kernels: bool = True, # Enable using SDPA from PyTroch Accelerated Transformers, make use Flash Attention and Xformer memory-efficient kernels
  32. **kwargs
  33. ):
  34. if prompt_file is not None:
  35. assert os.path.exists(
  36. prompt_file
  37. ), f"Provided Prompt file does not exist {prompt_file}"
  38. with open(prompt_file, "r") as f:
  39. user_prompt = f.read()
  40. else:
  41. print("No user prompt provided. Exiting.")
  42. sys.exit(1)
  43. # Set the seeds for reproducibility
  44. torch.cuda.manual_seed(seed)
  45. torch.manual_seed(seed)
  46. model = load_model(model_name, quantization)
  47. model.config.tp_size=1
  48. if peft_model:
  49. model = load_peft_model(model, peft_model)
  50. model.eval()
  51. if use_fast_kernels:
  52. """
  53. Setting 'use_fast_kernels' will enable
  54. using of Flash Attention or Xformer memory-efficient kernels
  55. based on the hardware being used. This would speed up inference when used for batched inputs.
  56. """
  57. try:
  58. from optimum.bettertransformer import BetterTransformer
  59. model = BetterTransformer.transform(model)
  60. except ImportError:
  61. print("Module 'optimum' not found. Please install 'optimum' it before proceeding.")
  62. tokenizer = AutoTokenizer.from_pretrained(model_name)
  63. safety_checker = get_safety_checker(enable_azure_content_safety,
  64. enable_sensitive_topics,
  65. enable_salesforce_content_safety,
  66. enable_llamaguard_content_safety,
  67. )
  68. # Safety check of the user prompt
  69. safety_results = [check(user_prompt) for check in safety_checker]
  70. are_safe = all([r[1] for r in safety_results])
  71. if are_safe:
  72. print("User prompt deemed safe.")
  73. print(f"User prompt:\n{user_prompt}")
  74. else:
  75. print("User prompt deemed unsafe.")
  76. for method, is_safe, report in safety_results:
  77. if not is_safe:
  78. print(method)
  79. print(report)
  80. print("Skipping the inference as the prompt is not safe.")
  81. sys.exit(1) # Exit the program with an error status
  82. batch = tokenizer(user_prompt, return_tensors="pt")
  83. batch = {k: v.to("cuda") for k, v in batch.items()}
  84. start = time.perf_counter()
  85. with torch.no_grad():
  86. outputs = model.generate(
  87. **batch,
  88. max_new_tokens=max_new_tokens,
  89. do_sample=do_sample,
  90. top_p=top_p,
  91. temperature=temperature,
  92. min_length=min_length,
  93. use_cache=use_cache,
  94. top_k=top_k,
  95. repetition_penalty=repetition_penalty,
  96. length_penalty=length_penalty,
  97. **kwargs
  98. )
  99. e2e_inference_time = (time.perf_counter()-start)*1000
  100. print(f"the inference time is {e2e_inference_time} ms")
  101. filling = tokenizer.batch_decode(outputs[:, batch["input_ids"].shape[1]:], skip_special_tokens=True)[0]
  102. # Safety check of the model output
  103. safety_results = [check(filling) for check in safety_checker]
  104. are_safe = all([r[1] for r in safety_results])
  105. if are_safe:
  106. print("User input and model output deemed safe.")
  107. print(user_prompt.replace("<FILL_ME>", filling))
  108. else:
  109. print("Model output deemed unsafe.")
  110. for method, is_safe, report in safety_results:
  111. if not is_safe:
  112. print(method)
  113. print(report)
  114. if __name__ == "__main__":
  115. fire.Fire(main)