update

open-compass · Mar 22, 2024 · 393da76 · 393da76
1 parent 22db9d9
commit 393da76
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/vlmeval/vlm/llava.py b/vlmeval/vlm/llava.py
@@ -201,4 +201,9 @@ def generate(self, image_path, prompt, dataset=None):
         inputs = self.processor(prompt_wtmpl, image, return_tensors='pt').to('cuda')
         output = self.model.generate(**inputs, **self.kwargs)
         answer = self.processor.decode(output[0], skip_special_token=True)
+        lt = len(prompt_wtmpl)
+        if prompt_wtmpl == answer[:lt]:
+            answer = answer[lt:]
+        answer = answer.split('</s>')[0]
+
         return answer