Share an easy-to-use Python + PyTorch + HuggingFace version of Grok-1

[binmakeswell]

https://huggingface.co/hpcai-tech/grok-1

[andy108369]

https://huggingface.co/hpcai-tech/grok-1

thanks for sharing it @binmakeswell !

it is working for me:

I have 8x H100

apt update && apt install -y python3-pip virtualenv git

mkdir grok-1
cd grok-1

virtualenv --python=python3 venv
source venv/bin/activate

pip install huggingface_hub[hf_transfer]

mkdir -p hpcaitech/grok-1
huggingface-cli download hpcai-tech/grok-1 --repo-type model --local-dir hpcaitech/grok-1 --local-dir-use-symlinks False

pip install torch transformers accelerate sentencepiece

wget https://github.com/xai-org/grok-1/raw/main/tokenizer.model

$ python3

import torch
import sentencepiece
from transformers import AutoModelForCausalLM

torch.set_default_dtype(torch.bfloat16)
model = AutoModelForCausalLM.from_pretrained(
    "hpcaitech/grok-1",
    trust_remote_code=True,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
sp = sentencepiece.SentencePieceProcessor(model_file="tokenizer.model")

text = "Who are you?"
input_ids = sp.encode(text)
input_ids = torch.tensor([input_ids]).cuda()
attention_mask = torch.ones_like(input_ids)

generate_kwargs = {
    'max_length': 100,  # Increase max_length to allow longer outputs
    'num_return_sequences': 1,  # Number of different sequences to generate
    # Add other parameters as needed
}

inputs = {
    "input_ids": input_ids,
    "attention_mask": attention_mask,
    **generate_kwargs,
}
outputs = model.generate(**inputs)

# Decode the output tensor to text
decoded_output = sp.decode(outputs[0].cpu().tolist())
print(decoded_output)

Result

>>> decoded_output = sp.decode(outputs[0].cpu().tolist())
>>> 
>>> print(decoded_output)

Who are you?

I am a 23 year old student from the Netherlands. I am currently in my last year of my bachelor in International Business and Management Studies. I am a very social person and I love to meet new people. I am also very interested in different cultures and I love to travel.

Why did you choose to do your internship at the Embassy of the Kingdom of the Netherlands in Bangkok?

I chose to do my internship at the Embassy of the Kingdom of the Netherlands
>>> 

It also looks like it doesn't require /dev/shm as opposed to the xai-org's grok-1 (probably JAX needs it?). Which is great since not all containers (K8s / Docker) have that.