Improve installation + code snippets

README.md

@@ -33,18 +33,12 @@ This repository contains [`meta-llama/Meta-Llama-3.1-405B-Instruct`](https://hug
 
 In order to use the current quantized model, support is offered for different solutions as `transformers`, `autoawq`, or `text-generation-inference`.
 
-### 🤗 transformers
+### 🤗 Transformers
 
-In order to run the inference with Llama 3.1 405B Instruct AWQ in INT4, both `torch` and `autoawq` need to be installed as:
+In order to run the inference with Llama 3.1 405B Instruct AWQ in INT4, you need to install the following packages:
 
 ```bash
-pip install "torch>=2.2.0,<2.3.0" autoawq --upgrade
-```
-
-Then, the latest version of `transformers` need to be installed, being 4.43.0 or higher, as:
-
-```bash
-pip install "transformers[accelerate]>=4.43.0" --upgrade
+pip install -q --upgrade transformers autoawq accelerate
 ```
 
 To run the inference on top of Llama 3.1 405B Instruct AWQ in INT4 precision, the AWQ model can be instantiated as any other causal language modeling model via `AutoModelForCausalLM` and run the inference normally.
@@ -54,15 +48,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model_id = "hugging-quants/Meta-Llama-3.1-405B-Instruct-AWQ-INT4"
-prompt = [
-  {"role": "system", "content": "You are a helpful assistant, that responds as a pirate."},
-  {"role": "user", "content": "What's Deep Learning?"},
-]
-
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-inputs = tokenizer.apply_chat_template(prompt, tokenize=True, add_generation_prompt=True, return_tensors="pt").cuda()
-
 model = AutoModelForCausalLM.from_pretrained(
   model_id,
   torch_dtype=torch.float16,
@@ -70,22 +56,28 @@ model = AutoModelForCausalLM.from_pretrained(
   device_map="auto",
 )
 
-outputs = model.generate(inputs, do_sample=True, max_new_tokens=256)
-print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
+prompt = [
+  {"role": "system", "content": "You are a helpful assistant, that responds as a pirate."},
+  {"role": "user", "content": "What's Deep Learning?"},
+]
+inputs = tokenizer.apply_chat_template(
+  prompt,
+  tokenize=True,
+  add_generation_prompt=True,
+  return_tensors="pt",
+  return_dict=True,
+).to("cuda")
+
+outputs = model.generate(**inputs, do_sample=True, max_new_tokens=256)
+print(tokenizer.batch_decode(outputs[:, inputs['input_ids'].shape[1]:], skip_special_tokens=True)[0])
 ```
 
 ### AutoAWQ
 
-In order to run the inference with Llama 3.1 405B Instruct AWQ in INT4, both `torch` and `autoawq` need to be installed as:
-
-```bash
-pip install "torch>=2.2.0,<2.3.0" autoawq --upgrade
-```
-
-Then, the latest version of `transformers` need to be installed, being 4.43.0 or higher, as:
+In order to run the inference with Llama 3.1 405B Instruct AWQ in INT4, you need to install the following packages:
 
 ```bash
-pip install "transformers[accelerate]>=4.43.0" --upgrade
+pip install -q --upgrade transformers autoawq accelerate
 ```
 
 Alternatively, one may want to run that via `AutoAWQ` even though it's built on top of 🤗 `transformers`, which is the recommended approach instead as described above.
@@ -96,11 +88,6 @@ from awq import AutoAWQForCausalLM
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model_id = "hugging-quants/Meta-Llama-3.1-405B-Instruct-AWQ-INT4"
-prompt = [
-  {"role": "system", "content": "You are a helpful assistant, that responds as a pirate."},
-  {"role": "user", "content": "What's Deep Learning?"},
-]
-
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoAWQForCausalLM.from_pretrained(
   model_id,
@@ -109,9 +96,20 @@ model = AutoAWQForCausalLM.from_pretrained(
   device_map="auto",
 )
 
-inputs = tokenizer.apply_chat_template(prompt, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to('cuda')
+prompt = [
+  {"role": "system", "content": "You are a helpful assistant, that responds as a pirate."},
+  {"role": "user", "content": "What's Deep Learning?"},
+]
+inputs = tokenizer.apply_chat_template(
+  prompt,
+  tokenize=True,
+  add_generation_prompt=True,
+  return_tensors="pt",
+  return_dict=True,
+).to("cuda")
+
 outputs = model.generate(**inputs, do_sample=True, max_new_tokens=256)
-print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
+print(tokenizer.batch_decode(outputs[:, inputs['input_ids'].shape[1]:], skip_special_tokens=True)[0])
 ```
 
 The AutoAWQ script has been adapted from [AutoAWQ/examples/generate.py](https://github.com/casper-hansen/AutoAWQ/blob/main/examples/generate.py).
@@ -125,21 +123,13 @@ Coming soon!
 > [!NOTE]
 > In order to quantize Llama 3.1 405B Instruct using AutoAWQ, you will need to use an instance with at least enough CPU RAM to fit the whole model i.e. ~800GiB, and an NVIDIA GPU with 80GiB of VRAM to quantize it.
 
-In order to quantize Llama 3.1 405B Instruct, first install `torch` and `autoawq` as follows:
-
-```bash
-pip install "torch>=2.2.0,<2.3.0" autoawq --upgrade
-```
-
-Otherwise the quantization may fail, since the AutoAWQ kernels are built with PyTorch 2.2.1, meaning that those will break with PyTorch 2.3.0.
-
-Then install the latest version of `transformers` as follows:
+In order to quantize Llama 3.1 405B Instruct, first install the following packages:
 
 ```bash
-pip install "transformers>=4.43.0" --upgrade
+pip install -q --upgrade transformers autoawq accelerate
 ```
 
-And then, run the following script, adapted from [`AutoAWQ/examples/quantize.py`](https://github.com/casper-hansen/AutoAWQ/blob/main/examples/quantize.py) as follows:
+Then run the following script, adapted from [`AutoAWQ/examples/quantize.py`](https://github.com/casper-hansen/AutoAWQ/blob/main/examples/quantize.py):
 
 ```python
 from awq import AutoAWQForCausalLM
@@ -156,9 +146,9 @@ quant_config = {
 
 # Load model
 model = AutoAWQForCausalLM.from_pretrained(
-  model_path, **{"low_cpu_mem_usage": True, "use_cache": False}
+  model_path, low_cpu_mem_usage=True, use_cache=False,
 )
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
 
 # Quantize
 model.quantize(tokenizer, quant_config=quant_config)