mirror of https://github.com/THUDM/ChatGLM-6B
Update loading instructions
parent
63b1e4e804
commit
5ade1e4055
|
@ -152,6 +152,7 @@ model = AutoModel.from_pretrained("THUDM/chatglm-6b", config=config, trust_remot
|
|||
prefix_state_dict = torch.load(os.path.join(CHECKPOINT_PATH, "pytorch_model.bin"))
|
||||
new_prefix_state_dict = {}
|
||||
for k, v in prefix_state_dict.items():
|
||||
if k.startswith("transformer.prefix_encoder."):
|
||||
new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
|
||||
model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
|
||||
```
|
||||
|
@ -160,7 +161,7 @@ model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
|
|||
(2) 如果需要加载的是旧 Checkpoint(包含 ChatGLM-6B 以及 PrefixEncoder 参数),或者进行的全参数微调,则直接加载整个 Checkpoint:
|
||||
|
||||
```python
|
||||
model = AutoModel.from_pretrained(CHECKPOINT_PATH, config=config, trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained(CHECKPOINT_PATH, trust_remote_code=True)
|
||||
```
|
||||
|
||||
之后根据需求可以进行量化,也可以直接使用:
|
||||
|
|
|
@ -118,6 +118,7 @@ def main():
|
|||
prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin"))
|
||||
new_prefix_state_dict = {}
|
||||
for k, v in prefix_state_dict.items():
|
||||
if k.startswith("transformer.prefix_encoder."):
|
||||
new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
|
||||
model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
|
||||
else:
|
||||
|
|
Loading…
Reference in New Issue