Update inference.py

This commit is contained in:
neukg 2023-06-29 19:57:14 +08:00 committed by GitHub
parent cdb808e7f6
commit 85ab09408f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions

View File

@ -1,7 +1,7 @@
from transformers import LlamaTokenizer, AutoModelForCausalLM, AutoConfig, GenerationConfig
import torch
ckpt_path = '/workspace/BELLE-train/Version_raw/'
ckpt_path = './ckpt_path'
load_type = torch.float16
device = torch.device(0)
tokenizer = LlamaTokenizer.from_pretrained(ckpt_path)