Update README.md

README.md

@@ -319,10 +319,10 @@ generate("I'm super excited about this Norwegian NORA model! Can it translate th
 ```
 
 _____
-## Example usage with low GPU usage
+## Example usage on a GPU with ~16GB VRAM
 Install bitsandbytes if you want to load in 8bit
 
-```python
+```bash
 pip install bitsandbytes
 pip install accelerate
 ```
@@ -334,13 +334,17 @@ import torch
 
 # First, we will have to import the tokenizer and the language model
 tokenizer = AutoTokenizer.from_pretrained("norallm/normistral-7b-warm")
-model = AutoModelForCausalLM.from_pretrained("norallm/normistral-7b-warm",
- device_map='auto',
- load_in_8bit=True,
- torch_dtype=torch.float16)
+
 # This setup needs about 8gb VRAM
-# Setting load_in_8bit = False, 15gb VRAM
-# Using torch.float32 and load_in_8bit = False, 21gb VRAM
+# Setting `load_in_8bit=False` -> 15gb VRAM
+# Using `torch.float32` and `load_in_8bit=False` -> 21gb VRAM
+model = AutoModelForCausalLM.from_pretrained(
+ "norallm/normistral-7b-warm",
+ device_map='auto',
+ load_in_8bit=True,
+ torch_dtype=torch.bfloat16
+)
+
 
 
 # Now we will define the zero-shot prompt template
@@ -362,5 +366,4 @@ def generate(text):
 
 # Now you can simply call the generate function with an English text you want to translate:
 generate("I'm super excited about this Norwegian NORA model! Can it translate these sentences?")
-# > this should output: 'Jeg er super spent på denne norske NORA modellen! Kan den oversette disse setningene?'
 ```