updated README.md

README.md

@@ -9,7 +9,7 @@ license_name: gemma-terms-of-use
 license_link: https://ai.google.dev/gemma/terms
 ---
 
-AWQ quantized version of gemma-2b model.
+AWQ quantized version of [google/gemma-2b](https://huggingface.co/google/gemma-2b).
 
 ---
 
@@ -61,22 +61,6 @@ In that repository, we provide:
 
 
 
-#### Running the model on a CPU
-
-
-```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-2b")
-
-input_text = "Write me a poem about Machine Learning."
-input_ids = tokenizer(input_text, return_tensors="pt")
-
-outputs = model.generate(**input_ids)
-print(tokenizer.decode(outputs[0]))
-```
-
 
 #### Running the model on a single / multi GPU
 
@@ -85,8 +69,8 @@ print(tokenizer.decode(outputs[0]))
 # pip install accelerate
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-AWQ")
+model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-AWQ", device_map="auto")
 
 input_text = "Write me a poem about Machine Learning."
 input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
@@ -104,62 +88,8 @@ print(tokenizer.decode(outputs[0]))
 # pip install accelerate
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto", torch_dtype=torch.float16)
-
-input_text = "Write me a poem about Machine Learning."
-input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
-
-outputs = model.generate(**input_ids)
-print(tokenizer.decode(outputs[0]))
-```
-
-* _Using `torch.bfloat16`_
-
-```python
-# pip install accelerate
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto", torch_dtype=torch.bfloat16)
-
-input_text = "Write me a poem about Machine Learning."
-input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
-
-outputs = model.generate(**input_ids)
-print(tokenizer.decode(outputs[0]))
-```
-
-#### Quantized Versions through `bitsandbytes`
-
-* _Using 8-bit precision (int8)_
-
-```python
-# pip install bitsandbytes accelerate
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-
-quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", quantization_config=quantization_config)
-
-input_text = "Write me a poem about Machine Learning."
-input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
-
-outputs = model.generate(**input_ids)
-print(tokenizer.decode(outputs[0]))
-```
-
-* _Using 4-bit precision_
-
-```python
-# pip install bitsandbytes accelerate
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-
-quantization_config = BitsAndBytesConfig(load_in_4bit=True)
-
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", quantization_config=quantization_config)
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-AWQ")
+model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-AWQ", device_map="auto", torch_dtype=torch.float16)
 
 input_text = "Write me a poem about Machine Learning."
 input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")