github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 10:34:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[coati] Fix LlamaCritic (#3475)

Browse Source 
* mv LlamaForCausalLM to LlamaModel

* rm unused imports

---------

Co-authored-by: gongenlei <gongenlei@baidu.com>
This commit is contained in:
gongenlei
2023-04-07 11:39:09 +08:00
committed by GitHub
parent 8f2c55f9c9
commit a7ca297281
1 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
applications/Chat/coati/models/llama/llama_critic.py
View File

@@ -1,8 +1,7 @@
from typing import Optional from typing import Optional
import torch
import torch.nn as nn import torch.nn as nn
from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM from transformers import LlamaConfig, LlamaModel
from ..base import Critic from ..base import Critic
@@ -28,11 +27,11 @@ class LlamaCritic(Critic):
**kwargs) -> None: **kwargs) -> None:
if pretrained is not None: if pretrained is not None:
model = LlamaForCausalLM.from_pretrained(pretrained) model = LlamaModel.from_pretrained(pretrained)
elif config is not None: elif config is not None:
model = LlamaForCausalLM(config) model = LlamaModel(config)
else: else:
model = LlamaForCausalLM(LlamaConfig()) model = LlamaModel(LlamaConfig())
if checkpoint: if checkpoint:
model.gradient_checkpointing_enable() model.gradient_checkpointing_enable()