github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix conversation: llama generates token not prob now (#3120)

Browse Source
This commit is contained in:
chenyu
2024-01-14 13:10:01 -05:00
committed by GitHub
parent 152ef7fc79
commit 79f4627fbc
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
examples/conversation.py
View File

@@ -79,8 +79,7 @@ def llama_generate(
outputted = llama.tokenizer.decode(toks)
init_length = len(outputted)
for _ in range(max_tokens):
probs_np = llama.model(Tensor([toks[start_pos:]]), start_pos, temperature).numpy()
token = int(np.random.choice(len(probs_np), p=probs_np))
token = llama.model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
start_pos = len(toks)
toks.append(token)