update README - usage of tokenizer

README.md

@@ -14,27 +14,29 @@ You could find the original weights released by [xAI](https://x.ai/blog) in [Hug
 
 We translated the original modeling written in JAX into PyTorch version, and converted the weights by mapping tensor files with parameter keys, de-quantizing the tensors with corresponding packed scales, and save to checkpoint file with torch APIs.
 
-The original tokenizer is supposed to be used (i.e. `tokenizer.model` in [GitHub Repository](https://github.com/xai-org/grok-1/tree/main)) with the torch-version model. 
+A transformers-compatible version of tokenizer is contributed by [Xenova](https://huggingface.co/Xenova) and [ArthurZ](https://huggingface.co/ArthurZ).
 
 ## Usage
 
 ```python
 import torch
-from transformers import AutoModelForCausalLM
-from sentencepiece import SentencePieceProcessor
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 torch.set_default_dtype(torch.bfloat16)
+
+tokenizer = AutoTokenizer.from_pretrained("hpcai-tech/grok-1", trust_remote_code=True)
+
 model = AutoModelForCausalLM.from_pretrained(
     "hpcai-tech/grok-1",
     trust_remote_code=True,
     device_map="auto",
     torch_dtype=torch.bfloat16,
 )
-sp = SentencePieceProcessor(model_file="tokenizer.model")
+model.eval()
 
 text = "Replace this with your text"
-input_ids = sp.encode(text)
-input_ids = torch.tensor([input_ids]).cuda()
+input_ids = tokenizer(text, return_tensors="pt").input_ids
+input_ids = input_ids.cuda()
 attention_mask = torch.ones_like(input_ids)
 generate_kwargs = {}  # Add any additional args if you want
 inputs = {
@@ -43,14 +45,7 @@ inputs = {
     **generate_kwargs,
 }
 outputs = model.generate(**inputs)
-```
-
-You could also use the transformers-compatible version of the tokenizer [Xenova/grok-1-tokenizer](https://huggingface.co/Xenova/grok-1-tokenizer) 
-```python
-from transformers import LlamaTokenizerFast
-
-tokenizer = LlamaTokenizerFast.from_pretrained('Xenova/grok-1-tokenizer')
-inputs = tokenizer('hello world')
+print(outputs)
 ```