423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-15 19:30:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix ChatGLM lm_head #494

Browse Source
This commit is contained in:
hiyouga
2023-08-14 14:14:48 +08:00
parent 20a29297b1
commit d019956808
3 changed files with 12 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/llmtuner/tuner/core/loader.py
View File

@@ -153,6 +153,10 @@ def load_model_and_tokenizer(
if "GenerationMixin" not in str(model.generate.__func__):
model.generate = MethodType(PreTrainedModel.generate, model)
# Fix LM head (for ChatGLM2)
if not hasattr(model, "lm_head"):
setattr(model, "lm_head", model.transformer.output_layer)
# Register auto class to save the custom code files.
if isinstance(config, PretrainedConfig) and "AutoConfig" in getattr(config, "auto_map", {}):
config.__class__.register_for_auto_class()