root
/
TechGPT-1.0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
inference.py
View File

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