Update README.md

README.md

@@ -18,6 +18,39 @@ The codebase for APUS-xDAN-4.0-MOE is integrated into the latest Hugging Face tr
 
 Copy code
 Usage llama.cpp
+
+Usage
+
+"""
+import torch
+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,
+)
+model.eval()
+
+text = "Replace this with your text"
+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 = {
+    "input_ids": input_ids,
+    "attention_mask": attention_mask,
+    **generate_kwargs,
+}
+outputs = model.generate(**inputs)
+print(outputs)
+"""
+
 License
 
 APUS-xDAN-4.0-MOE is distributed under the LLAMA 2 Community License, Copyright (c) Meta Platforms, Inc. All Rights Reserved.