Update README.md

README.md

@@ -53,7 +53,7 @@ messages = [
 
 prompts = tokenizer.apply_chat_template(messages, tokenize=False)
 
-llm = LLM(model=model_id)
+llm = LLM(model=model_id, tensor_parallel_size=1)
 
 outputs = llm.generate(prompts, sampling_params)
 
@@ -124,9 +124,9 @@ num_samples = 512
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-preprocess_fn = lambda example: {"text": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n{output}".format_map(example)
+preprocess_fn = lambda example: {"text": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n{text}".format_map(example)}
 
-dataset_name = "garage-bAInd/Open-Platypus"
+dataset_name = "neuralmagic/LLM_compression_calibration"
 dataset = load_dataset(dataset_name, split="train")
 ds = dataset.shuffle().select(range(num_samples))
 ds = ds.map(preprocess_fn)