alihan/llama-chat-llm
1
0
Fork 0
mirror of https://github.com/randaller/llama-chat.git synced 2023-09-17 22:41:47 +03:00
Code Issues Packages Projects Releases Wiki Activity

Update hf-inference-cuda-example.py

Browse Source
This commit is contained in:
randaller
2023-03-19 19:35:20 +03:00
committed by GitHub
parent 9560b01fec
commit 4ef8416f5f
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
hf-inference-cuda-example.py
View File

@@ -22,4 +22,4 @@ device_map = infer_auto_device_map(model, max_memory={0: "6GiB", "cpu": "128GiB"
print(device_map)
batch = tokenizer("The highest mountain in China is ", return_tensors="pt")
print(tokenizer.decode(model.generate(batch["input_ids"].cpu(), do_sample=True, top_k=50, max_length=100, top_p=0.95, temperature=1.0)[0]))
print(tokenizer.decode(model.generate(batch["input_ids"].cuda(), do_sample=True, top_k=50, max_length=100, top_p=0.95, temperature=1.0)[0]))