Offline Inference Encoder Decoder

Source vllm-project/vllm.

  1.  1'''
  2.  2Demonstrate prompting of text-to-text
  3.  3encoder/decoder models, specifically BART
  4.  4'''
  5.  5
  6.  6from vllm import LLM, SamplingParams
  7.  7from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
  8.  8                         TokensPrompt, zip_enc_dec_prompts)
  9.  9
  10. 10dtype = "float"
  11. 11
  12. 12# Create a BART encoder/decoder model instance
  13. 13llm = LLM(
  14. 14    model="facebook/bart-large-cnn",
  15. 15    dtype=dtype,
  16. 16)
  17. 17
  18. 18# Get BART tokenizer
  19. 19tokenizer = llm.llm_engine.get_tokenizer_group()
  20. 20
  21. 21# Test prompts
  22. 22#
  23. 23# This section shows all of the valid ways to prompt an
  24. 24# encoder/decoder model.
  25. 25#
  26. 26# - Helpers for building prompts
  27. 27text_prompt_raw = "Hello, my name is"
  28. 28text_prompt = TextPrompt(prompt="The president of the United States is")
  29. 29tokens_prompt = TokensPrompt(prompt_token_ids=tokenizer.encode(
  30. 30    prompt="The capital of France is"))
  31. 31# - Pass a single prompt to encoder/decoder model
  32. 32#   (implicitly encoder input prompt);
  33. 33#   decoder input prompt is assumed to be None
  34. 34
  35. 35single_text_prompt_raw = text_prompt_raw  # Pass a string directly
  36. 36single_text_prompt = text_prompt  # Pass a TextPrompt
  37. 37single_tokens_prompt = tokens_prompt  # Pass a TokensPrompt
  38. 38
  39. 39# - Pass explicit encoder and decoder input prompts within one data structure.
  40. 40#   Encoder and decoder prompts can both independently be text or tokens, with
  41. 41#   no requirement that they be the same prompt type. Some example prompt-type
  42. 42#   combinations are shown below, note that these are not exhaustive.
  43. 43
  44. 44enc_dec_prompt1 = ExplicitEncoderDecoderPrompt(
  45. 45    # Pass encoder prompt string directly, &
  46. 46    # pass decoder prompt tokens
  47. 47    encoder_prompt=single_text_prompt_raw,
  48. 48    decoder_prompt=single_tokens_prompt,
  49. 49)
  50. 50enc_dec_prompt2 = ExplicitEncoderDecoderPrompt(
  51. 51    # Pass TextPrompt to encoder, and
  52. 52    # pass decoder prompt string directly
  53. 53    encoder_prompt=single_text_prompt,
  54. 54    decoder_prompt=single_text_prompt_raw,
  55. 55)
  56. 56enc_dec_prompt3 = ExplicitEncoderDecoderPrompt(
  57. 57    # Pass encoder prompt tokens directly, and
  58. 58    # pass TextPrompt to decoder
  59. 59    encoder_prompt=single_tokens_prompt,
  60. 60    decoder_prompt=single_text_prompt,
  61. 61)
  62. 62
  63. 63# - Finally, here's a useful helper function for zipping encoder and
  64. 64#   decoder prompts together into a list of ExplicitEncoderDecoderPrompt
  65. 65#   instances
  66. 66zipped_prompt_list = zip_enc_dec_prompts(
  67. 67    ['An encoder prompt', 'Another encoder prompt'],
  68. 68    ['A decoder prompt', 'Another decoder prompt'])
  69. 69
  70. 70# - Let's put all of the above example prompts together into one list
  71. 71#   which we will pass to the encoder/decoder LLM.
  72. 72prompts = [
  73. 73    single_text_prompt_raw, single_text_prompt, single_tokens_prompt,
  74. 74    enc_dec_prompt1, enc_dec_prompt2, enc_dec_prompt3
  75. 75] + zipped_prompt_list
  76. 76
  77. 77print(prompts)
  78. 78
  79. 79# Create a sampling params object.
  80. 80sampling_params = SamplingParams(
  81. 81    temperature=0,
  82. 82    top_p=1.0,
  83. 83    min_tokens=0,
  84. 84    max_tokens=20,
  85. 85)
  86. 86
  87. 87# Generate output tokens from the prompts. The output is a list of
  88. 88# RequestOutput objects that contain the prompt, generated
  89. 89# text, and other information.
  90. 90outputs = llm.generate(prompts, sampling_params)
  91. 91
  92. 92# Print the outputs.
  93. 93for output in outputs:
  94. 94    prompt = output.prompt
  95. 95    encoder_prompt = output.encoder_prompt
  96. 96    generated_text = output.outputs[0].text
  97. 97    print(f"Encoder prompt: {encoder_prompt!r}, "
  98. 98          f"Decoder prompt: {prompt!r}, "
  99. 99          f"Generated text: {generated_text!r}")