update

vlmeval/vlm/llava.py

@@ -132,7 +132,10 @@ def __init__(self, model_pth='llava-hf/llava-v1.6-vicuna-7b-hf', **kwargs):
         assert version_cmp(transformers.__version__, '4.39.0', 'ge')
         from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
         self.model_pth = model_pth
-        self.processor = LlavaNextProcessor.from_pretrained(self.model_pth)
+        if '34b' in model_pth.lower():
+            self.processor = LlavaNextProcessor.from_pretrained(self.model_pth, use_fast=False)
+        else:
+            self.processor = LlavaNextProcessor.from_pretrained(self.model_pth)
         model = LlavaNextForConditionalGeneration.from_pretrained(
             self.model_pth, torch_dtype=torch.float16, low_cpu_mem_usage=True)
         model = model.eval()
@@ -202,15 +205,15 @@ def generate(self, image_path, prompt, dataset=None):
         output = self.model.generate(**inputs, **self.kwargs)
         answer = self.processor.decode(output[0], skip_special_token=True)
         if '<s>' in answer:
-            answer = answer.replace('<s>', '')
-            answer = answer.strip()
-
-        lt = len(prompt_wtmpl)
-        if prompt_wtmpl == answer[:lt]:
-            answer = answer[lt:]
+            answer = answer.replace('<s>', '').strip()
+        if '[/INST]' in answer:
+            answer = answer.split('[/INST]')[1].strip()
+        elif 'ASSISTANT:' in answer:
+            answer = answer.split('ASSISTANT:')[1].strip()
+        elif '<|im_start|>assistant\n' in answer:
+            answer = answer.split('<|im_start|>assistant\n')[1].strip()
 
         if '</s>' in answer:
-            answer = answer.split('</s>')[0]
-            answer = answer.strip()
+            answer = answer.split('</s>')[0].strip()
 
         return answer