Update README.md

README.md

@@ -42,27 +42,24 @@ Only the weights and activations of the linear operators within transformers blo
 This model can be deployed efficiently using the [vLLM](https://docs.vllm.ai/en/latest/) backend, as shown in the example below.
 
 ```python
-from vllm import LLM, SamplingParams
 from transformers import AutoTokenizer
+from vllm import LLM, SamplingParams
 
-model_id = "neuralmagic/DeepSeek-Coder-V2-Base-FP8"
-
-sampling_params = SamplingParams(temperature=0.6, top_p=0.9, max_tokens=256)
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+max_model_len, tp_size = 4096, 4
+model_name = "neuralmagic/DeepSeek-Coder-V2-Base-FP8"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+llm = LLM(model=model_name, tensor_parallel_size=tp_size, max_model_len=max_model_len, trust_remote_code=True, enforce_eager=True)
+sampling_params = SamplingParams(temperature=0.3, max_tokens=256, stop_token_ids=[tokenizer.eos_token_id])
 
-messages = [
-    {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
-    {"role": "user", "content": "Who are you?"},
+messages_list = [
+    [{"role": "user", "content": "Who are you? Please respond in pirate speak!"}],
 ]
 
-prompts = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-llm = LLM(model=model_id, trust_remote_code=True, max_model_len=4096, tensor_parallel_size=4)
+prompt_token_ids = [tokenizer.apply_chat_template(messages, add_generation_prompt=True) for messages in messages_list]
 
-outputs = llm.generate(prompts, sampling_params)
+outputs = llm.generate(prompt_token_ids=prompt_token_ids, sampling_params=sampling_params)
 
-generated_text = outputs[0].outputs[0].text
+generated_text = [output.outputs[0].text for output in outputs]
 print(generated_text)
 ```