Improve installation + code snippets

README.md

@@ -33,20 +33,12 @@ This repository contains [`meta-llama/Meta-Llama-3.1-70B-Instruct`](https://hugg
 
 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 70B Instruct AWQ in INT4, both `torch` and `autoawq` need to be installed as:
+In order to run the inference with Llama 3.1 70B Instruct AWQ in INT4, you need to install the following packages:
 
 ```bash
-pip install "torch>=2.2.0,<2.3.0" autoawq --upgrade
-```
-
-Otherwise, running the inference may fail, since the AutoAWQ kernels are built with PyTorch 2.2.1, meaning that those will break with PyTorch 2.3.0.
-
-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 70B 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.
@@ -56,13 +48,18 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model_id = "hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+  model_id,
+  torch_dtype=torch.float16,
+  low_cpu_mem_usage=True,
+  device_map="auto",
+)
+
 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,
@@ -71,48 +68,38 @@ inputs = tokenizer.apply_chat_template(
   return_dict=True,
 ).to("cuda")
 
-model = AutoModelForCausalLM.from_pretrained(
-  model_id,
-  torch_dtype=torch.float16,
-  low_cpu_mem_usage=True,
-  device_map="auto",
-)
-
 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])
 ```
 
 ### AutoAWQ
 
-In order to run the inference with Llama 3.1 70B 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
-```
-
-Otherwise, running the inference may fail, since the AutoAWQ kernels are built with PyTorch 2.2.1, meaning that those will break with PyTorch 2.3.0.
-
-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 70B 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.
 
 ```python
 import torch
-from autoawq import AutoAWQForCausalLM
+from awq import AutoAWQForCausalLM
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model_id = "hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoAWQForCausalLM.from_pretrained(
+  model_id,
+  torch_dtype=torch.float16,
+  low_cpu_mem_usage=True,
+  device_map="auto",
+)
+
 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,
@@ -121,15 +108,8 @@ inputs = tokenizer.apply_chat_template(
   return_dict=True,
 ).to("cuda")
 
-model = AutoAWQForCausalLM.from_pretrained(
-  model_id,
-  torch_dtype=torch.float16,
-  low_cpu_mem_usage=True,
-  device_map="auto",
-)
-
 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).
@@ -143,21 +123,13 @@ Coming soon!
 > [!NOTE]
 > In order to quantize Llama 3.1 70B Instruct using AutoAWQ, you will need to use an instance with at least enough CPU RAM to fit the whole model i.e. ~140GiB, and an NVIDIA GPU with 40GiB of VRAM to quantize it.
 
-In order to quantize Llama 3.1 70B 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 70B 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
@@ -174,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)