OpenCompass/opencompass/multimodal/models/visualglm/post_processor.py
Yike Yuan 97fdc51102
[Fix] Fix performance issue of visualglm. (#424)
* [Fix] Visualglm performance fixed.

* [Fix] Hide ckpt path.
2023-09-21 19:54:23 +08:00

30 lines
757 B
Python

from typing import Any
import torch
class VisualGLMBasePostProcessor:
"""Base post processor for VisualGLM."""
def __init__(self) -> None:
pass
def __call__(self, output_token: torch.tensor, tokenizer: Any) -> str:
return tokenizer.decode(output_token)
class VisualGLMVSRPostProcessor(VisualGLMBasePostProcessor):
"""VSR post processor for VisualGLM."""
def __init__(self) -> None:
super().__init__()
def __call__(self, output_token: torch.tensor, tokenizer: Any) -> str:
output_text = tokenizer.decode(output_token)
if 'yes' in output_text.lower():
return 'yes'
elif 'no' in output_text.lower():
return 'no'
else:
return 'unknown'