From a80a6066853e3f9615c7bba6fbe8d6e3a7534892 Mon Sep 17 00:00:00 2001 From: randaller Date: Sun, 19 Mar 2023 12:50:59 +0300 Subject: [PATCH] Update hf-inference-example.py --- hf-inference-example.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/hf-inference-example.py b/hf-inference-example.py index af59c80..5ce2558 100644 --- a/hf-inference-example.py +++ b/hf-inference-example.py @@ -5,9 +5,11 @@ MODEL = 'decapoda-research/llama-7b-hf' # MODEL = 'decapoda-research/llama-30b-hf' # MODEL = 'decapoda-research/llama-65b-hf' +# MODEL = './trained' + tokenizer = llamahf.LLaMATokenizer.from_pretrained(MODEL) model = llamahf.LLaMAForCausalLM.from_pretrained(MODEL, low_cpu_mem_usage=True) model.to('cpu') batch = tokenizer("The highest mountain in China is ", return_tensors="pt") -print(tokenizer.decode(model.generate(batch["input_ids"].cpu(), max_length=256)[0])) +print(tokenizer.decode(model.generate(batch["input_ids"].cpu(), max_length=100)[0]))