File size: 2,228 Bytes
cb2428f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
from dataclasses import asdict
from typing import Any, Dict, List, Union
import torch.nn as nn
from evalscope.models.custom import CustomModel
from transformers import PreTrainedModel
from ..infer import PtEngine, RequestConfig
from ..template import InferRequest
class EvalModel(CustomModel):
def __init__(self, model: Union[PreTrainedModel, nn.Module], template, max_batch_size, model_name: str,
**kwargs) -> None:
super().__init__(config={'model_id': model_name}, **kwargs)
self.model_name = model_name
self.model = model
self.template = template
self.engine = PtEngine.from_model_template(model, template, max_batch_size=max_batch_size)
def predict(self, prompts: List[dict], **kwargs) -> List[Dict[str, Any]]:
# use origin inputs
infer_requests = self.prepare_inputs(kwargs.get('origin_inputs', prompts))
infer_cfg = kwargs['infer_cfg'].copy()
generation_config = RequestConfig(**infer_cfg)
response = self.engine.infer(infer_requests=infer_requests, request_config=generation_config, use_tqdm=False)
dict_response = [asdict(item) for item in response]
return dict_response
def prepare_inputs(self, prompts: Union[List[dict], List[str]]) -> List[InferRequest]:
infer_requests = []
for input_item in prompts:
if isinstance(input_item, str):
query = input_item
system_prompt = None
else:
data: list = input_item['data']
if isinstance(data[0], tuple): # for truthful_qa and hellaswag
query = '\n'.join(''.join(item) for item in data)
system_prompt = input_item.get('system_prompt', None)
else:
query = data[0]
system_prompt = input_item.get('system_prompt', None)
# prepare messages
messages = []
if system_prompt:
messages.append({'role': 'system', 'content': system_prompt})
messages.append({'role': 'user', 'content': query})
infer_requests.append(InferRequest(messages=messages))
return infer_requests
|