fix bug

mrG7 · Jan 3, 2024 · 6b91dba · 6b91dba
1 parent b939855
commit 6b91dba
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/vlmeval/vlm/cogvlm.py b/vlmeval/vlm/cogvlm.py
@@ -4,10 +4,10 @@
 import os
 import os.path as osp
 from ..smp import *
-from ..utils import DATASET_TYPE
-from transformers import AutoModelForCausalLM, AutoTokenizer,LlamaTokenizer
+from ..utils import DATASET_TYPE, CustomPrompt
+from transformers import AutoModelForCausalLM, LlamaTokenizer
 
-class CogVlm:
+class CogVlm(CustomPrompt):
 
     INSTALL_REQ = True
 

diff --git a/vlmeval/vlm/sharedcaptioner.py b/vlmeval/vlm/sharedcaptioner.py
@@ -3,9 +3,9 @@
 import warnings
 import os.path as osp
 from ..smp import *
-from ..utils import DATASET_TYPE
+from ..utils import DATASET_TYPE, CustomPrompt
 
-class SharedCaptioner:
+class SharedCaptioner(CustomPrompt):
 
     INSTALL_REQ = False