Update README.md

README.md

@@ -1,41 +1,83 @@
 ---
-base_model: GeneZC/MiniChat-2-3B
+base_model: nm-testing/MiniChat-2-3B-pruned2.4
 inference: True
 model_type: Llama
+tags:
+- nm-vllm
+- sparse
 ---
-#  MiniChat-2-3B
-This repo contains pruned model files for [MiniChat-2-3B](https://huggingface.co/GeneZC/MiniChat-2-3B).
+## MiniChat-2-3B-pruned2.4
+This repo contains model files for [MiniChat-2-3B-pruned2.4](https://huggingface.co/GeneZC/MiniChat-2-3B) optimized for [NM-vLLM](https://github.com/neuralmagic/nm-vllm), a high-throughput serving engine for compressed LLMs.
 
 This model was pruned with [SparseGPT](https://arxiv.org/abs/2301.00774), using [SparseML](https://github.com/neuralmagic/sparseml).
+
+## Inference
+Install [NM-vLLM](https://github.com/neuralmagic/nm-vllm) for fast inference and low memory-usage: 
+```bash
+pip install nm-vllm[sparse]
+```
+Run in a Python pipeline for local inference:
 ```python
+from vllm import LLM, SamplingParams
 
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+model = LLM("nm-testing/MiniChat-2-3B-pruned2.4", sparsity="sparse_w16a16")
 prompt = "How to make banana bread?"
 formatted_prompt =  f"<s> [|User|]\n{prompt}</s>[|Assistant|]\n"
-model_id = "nm-testing/MiniChat-2-3B-pruned50-24"
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-inputs = tokenizer(formatted_prompt, return_tensors="pt")
-outputs = model.generate(**inputs, max_new_tokens=200)
-print(tokenizer.batch_decode(outputs)[0])
 
+sampling_params = SamplingParams(max_tokens=100,temperature=0,repetition_penalty=1.3)
+outputs = model.generate(formatted_prompt, sampling_params=sampling_params)
+print(outputs[0].outputs[0].text)
 """
-<s><s> [|User|]
-How to make banana bread?</s>[|Assistant|]
-To make banana bread, follow these steps:
-
-1. Start by preparing the ingredients. You will need banana bread mix, flour, water, and salt.
-2. Mix the ingredients together and mix the mixture thoroughly.
-3. Pour the mixture into a pan to cook.
-4. Cook the mixture until it is cooked.
-5. Once the bread is cooked, you can use it as a base for making banana bread.
-6. Add the banana bread mix to the pan and mix it thoroughly.
-7. Pour the mixture into a pan to cook.
-8. Cook the mixture until it is cooked.
-9. Once the bread is cooked, you can use it as a base for making banana bread.
-10. Add the banana bread mix to the pan and mix it thoroughly.
-11. Pour the mixture into a pan to cook.
-12.
+Answer: Create a recipe for making banana bread using ingredients like flour, water and sugar.  Explain the process of mixing these materials together until they form an unpleasant mixture that can be used in cooking methods such as baking or boiling processes. Describe how you would create this dough by adding it into your kitchen's oven-based environment while describing its properties during each stage before creating them on topical forms. You will also describe what
+
 """
-```
+```
+
+## Prompt template
+
+```
+### User:
+{prompt}
+### Assistant:
+
+```
+
+## Sparsification
+For details on how this model was sparsified, see the `recipe.yaml` in this repo and follow the instructions below.
+
+Install [SparseML](https://github.com/neuralmagic/sparseml):
+```bash
+git clone https://github.com/neuralmagic/sparseml
+pip install -e "sparseml[transformers]"
+```
+
+Replace the recipe as you like and run this one-shot compression script to apply SparseGPT:
+```python
+import sparseml.transformers
+
+original_model_name = "nm-testing/MiniChat-2-3B-pruned2.4"
+calibration_dataset = "open_platypus"
+output_directory = "output/"
+
+recipe = """
+test_stage:
+  obcq_modifiers:
+    SparseGPTModifier:
+      sparsity: 0.5
+      sequential_update: true
+      mask_structure: '2:4'
+      targets: ['re:model.layers.\d*$']
+"""
+
+# Apply SparseGPT to the model
+sparseml.transformers.oneshot(
+    model=original_model_name,
+    dataset=calibration_dataset,
+    recipe=recipe,
+    output_dir=output_directory,
+)
+```
+
+## Slack
+
+For further support, and discussions on these models and AI in general, join [Neural Magic's Slack Community](https://join.slack.com/t/discuss-neuralmagic/shared_invite/zt-q1a1cnvo-YBoICSIw3L1dmQpjBeDurQ)