Update README.md

README.md

@@ -46,7 +46,7 @@ Find below some example scripts on how to use the model in `transformers` (Make
 
 ## Using the Pytorch model
 
-### Running the model on a CPU
+This checkpoint will only run on a GPU device with `bitsandbytes` installed. See below for more details on how to load it
 
 <details>
 <summary> Click to expand </summary>
@@ -54,8 +54,8 @@ Find below some example scripts on how to use the model in `transformers` (Make
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-mamba-7b")
-model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b")
+tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-mamba-7b-4bit")
+model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b-4bit")
 
 input_text = "Question: How many hours in one day? Answer: "
 input_ids = tokenizer(input_text, return_tensors="pt").input_ids
@@ -66,89 +66,21 @@ print(tokenizer.decode(outputs[0]))
 
 </details>
 
-### Running the model on a GPU
+You can also dequantize the model with `model.dequantize()` method:
 
-<details>
-<summary> Click to expand </summary>
-
-```python
-# pip install accelerate
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-mamba-7b")
-model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b", device_map="auto")
-
-input_text = "Question: How many hours in one day? Answer: "
-input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
-
-outputs = model.generate(input_ids)
-print(tokenizer.decode(outputs[0]))
-```
-
-</details>
-
-### Running the model on a GPU using `torch.compile`
 
 <details>
 <summary> Click to expand </summary>
 
 ```python
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-mamba-7b")
-model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b", torch_dtype=torch.bfloat16).to(0)
-
-model = torch.compile(model)
+tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-mamba-7b-4bit")
+model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b-4bit")
+model = model.dequantize()
 
 input_text = "Question: How many hours in one day? Answer: "
-input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
-
-outputs = model.generate(input_ids)
-print(tokenizer.decode(outputs[0]))
-```
-
-</details>
-
-
-### Running the model on a GPU using different precisions
-
-#### FP16
-
-<details>
-<summary> Click to expand </summary>
-
-```python
-# pip install accelerate
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-mamba-7b")
-model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b", device_map="auto", torch_dtype=torch.float16)
-
-input_text = "Question: How many hours in one day? Answer: "
-input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
-
-outputs = model.generate(input_ids)
-print(tokenizer.decode(outputs[0]))
-```
-
-</details>
-
-#### 4-bit
-
-<details>
-<summary> Click to expand </summary>
-
-```python
-# pip install bitsandbytes accelerate
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-
-tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-mamba-7b")
-model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b", device_map="auto", quantization_config=BitsAndBytesConfig(load_in_4bit=True))
-
-input_text = "Question: How many hours in one day? Answer: "
-input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids
 
 outputs = model.generate(input_ids)
 print(tokenizer.decode(outputs[0]))
@@ -157,7 +89,6 @@ print(tokenizer.decode(outputs[0]))
 </details>
 
 
-
 # Training Details
 
 ## Training Data