Witllm/chatglm/train.py

49 lines
1.5 KiB
Python
Raw Normal View History

2024-01-04 19:56:30 +08:00
import sys
sys.path.append("..")
2023-12-26 14:08:02 +08:00
import json
import torch
2024-01-04 19:56:30 +08:00
from modeling_chatglm import ChatGLMForConditionalGeneration
from tokenization_chatglm import ChatGLMTokenizer
from modelscope import snapshot_download
2023-12-26 14:08:02 +08:00
2024-01-03 20:26:26 +08:00
from tools import show
from transformers import AutoConfig
seed = 4321
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
2024-01-04 19:56:30 +08:00
pretrained_model_name_or_path = snapshot_download("ZhipuAI/chatglm3-6b")
2024-01-03 20:26:26 +08:00
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)
2023-12-26 14:08:02 +08:00
2024-01-04 19:56:30 +08:00
tokenizer_config_file = "./tokenizer_config.json"
2023-12-26 14:08:02 +08:00
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
2024-01-04 19:56:30 +08:00
init_kwargs["vocab_file"] = "./tokenizer.model"
2023-12-26 14:08:02 +08:00
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)
2024-01-03 20:26:26 +08:00
glm = glm.from_pretrained(pretrained_model_name_or_path).half().cuda()
query = "你好"
response = glm.backward(tokenizer, query)