46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
import json
|
|
import torch
|
|
|
|
from chatglm import ChatGLMForConditionalGeneration
|
|
from chatglm import ChatGLMTokenizer
|
|
|
|
from tools import show
|
|
|
|
from transformers import AutoConfig
|
|
|
|
seed = 4321
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
|
|
pretrained_model_name_or_path = "../ZhipuAI/chatglm3-6b"
|
|
config, kwargs = AutoConfig.from_pretrained(
|
|
pretrained_model_name_or_path,
|
|
return_unused_kwargs=True,
|
|
trust_remote_code=True,
|
|
code_revision=None,
|
|
_commit_hash=None,
|
|
)
|
|
glm = ChatGLMForConditionalGeneration(config)
|
|
|
|
|
|
tokenizer_config_file = "./chatglm/tokenizer_config.json"
|
|
if tokenizer_config_file is not None:
|
|
with open(tokenizer_config_file, encoding="utf-8") as tokenizer_config_handle:
|
|
init_kwargs = json.load(tokenizer_config_handle)
|
|
init_kwargs.pop("tokenizer_class", None)
|
|
init_kwargs.pop("tokenizer_file", None)
|
|
saved_init_inputs = init_kwargs.pop("init_inputs", ())
|
|
init_inputs = saved_init_inputs
|
|
init_kwargs["vocab_file"] = "./chatglm/tokenizer.model"
|
|
init_kwargs["added_tokens_file"] = None
|
|
init_kwargs["special_tokens_map_file"] = None
|
|
init_kwargs["tokenizer_file"] = None
|
|
init_kwargs["name_or_path"] = pretrained_model_name_or_path
|
|
tokenizer = ChatGLMTokenizer(*init_inputs, **init_kwargs)
|
|
|
|
|
|
glm = glm.from_pretrained(pretrained_model_name_or_path).half().cuda()
|
|
query = "你好"
|
|
response = glm.backward(tokenizer, query)
|