github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-10-01 06:25:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[example] update Grok-1 inference (#5495)

Browse Source 
* revise grok-1 example

* remove unused arg in scripts

* prevent re-installing torch

* update readme

* revert modifying colossalai requirements

* add perf

* trivial

* add tokenizer url
This commit is contained in:
Yuanheng Zhao
2024-03-24 20:24:11 +08:00
committed by GitHub
parent 6df844b8c4
commit 5fcd7795cd
7 changed files with 69 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
examples/language/grok-1/utils.py
View File

@@ -20,9 +20,9 @@ def print_output(text, output):
@torch.no_grad()
def inference(model, sp, text, **generate_kwargs):
input_ids = sp.encode(text)
input_ids = torch.tensor([input_ids]).cuda()
def inference(model, tokenizer, text, **generate_kwargs):
input_ids = tokenizer(text, return_tensors="pt").input_ids
input_ids = input_ids.cuda()
attention_mask = torch.ones_like(input_ids)
inputs = {
"input_ids": input_ids,