| | import json |
| | from concurrent.futures import ThreadPoolExecutor |
| | from typing import Dict, List, Optional |
| |
|
| | import numpy as np |
| | import requests |
| |
|
| | from opencompass.registry import MODELS |
| | from opencompass.utils.logging import get_logger |
| |
|
| | from .base_api import BaseAPIModel |
| |
|
| |
|
| | @MODELS.register_module() |
| | class LightllmAPI(BaseAPIModel): |
| |
|
| | is_api: bool = True |
| |
|
| | def __init__( |
| | self, |
| | path: str = 'LightllmAPI', |
| | url: str = 'http://localhost:8080/generate', |
| | input_format: str = '<input_text_to_replace>', |
| | max_seq_len: int = 2048, |
| | meta_template: Optional[Dict] = None, |
| | retry: int = 2, |
| | generation_kwargs: Optional[Dict] = dict(), |
| | ): |
| |
|
| | super().__init__(path=path, |
| | max_seq_len=max_seq_len, |
| | meta_template=meta_template, |
| | retry=retry, |
| | generation_kwargs=generation_kwargs) |
| | self.logger = get_logger() |
| | self.url = url |
| | self.input_format = input_format |
| | self.generation_kwargs = generation_kwargs |
| | self.max_out_len = self.generation_kwargs.get('max_new_tokens', 1024) |
| |
|
| | def generate(self, inputs: List[str], max_out_len: int, |
| | **kwargs) -> List[str]: |
| | """Generate results given a list of inputs. |
| | |
| | Args: |
| | inputs (List[str]): A list of strings or PromptDicts. |
| | The PromptDict should be organized in OpenCompass' |
| | API format. |
| | max_out_len (int): The maximum length of the output. |
| | |
| | Returns: |
| | List[str]: A list of generated strings. |
| | """ |
| |
|
| | with ThreadPoolExecutor() as executor: |
| | results = list( |
| | executor.map(self._generate, inputs, |
| | [self.max_out_len] * len(inputs))) |
| | return results |
| |
|
| | def _generate(self, input: str, max_out_len: int) -> str: |
| | max_num_retries = 0 |
| | while max_num_retries < self.retry: |
| | self.wait() |
| | header = {'content-type': 'application/json'} |
| | try: |
| | input = self.input_format.replace('<input_text_to_replace>', |
| | input) |
| | data = dict(inputs=input, parameters=self.generation_kwargs) |
| | raw_response = requests.post(self.url, |
| | headers=header, |
| | data=json.dumps(data)) |
| | except requests.ConnectionError: |
| | self.logger.error('Got connection error, retrying...') |
| | continue |
| | try: |
| | response = raw_response.json() |
| | generated_text = response['generated_text'] |
| | if isinstance(generated_text, list): |
| | generated_text = generated_text[0] |
| | return generated_text |
| | except requests.JSONDecodeError: |
| | self.logger.error('JsonDecode error, got', |
| | str(raw_response.content)) |
| | except KeyError: |
| | self.logger.error(f'KeyError. Response: {str(response)}') |
| | max_num_retries += 1 |
| |
|
| | raise RuntimeError('Calling LightllmAPI failed after retrying for ' |
| | f'{max_num_retries} times. Check the logs for ' |
| | 'details.') |
| |
|
| | def get_ppl(self, inputs: List[str], max_out_len: int, |
| | **kwargs) -> List[float]: |
| | """Generate results given a list of inputs. |
| | |
| | Args: |
| | inputs (List[str]): A list of strings or PromptDicts. |
| | The PromptDict should be organized in OpenCompass' |
| | API format. |
| | max_out_len (int): The maximum length of the output. |
| | |
| | Returns: |
| | List[str]: A list of generated strings. |
| | """ |
| |
|
| | with ThreadPoolExecutor() as executor: |
| | results = list( |
| | executor.map(self._get_ppl, inputs, |
| | [self.max_out_len] * len(inputs))) |
| | return np.array(results) |
| |
|
| | def _get_ppl(self, input: str, max_out_len: int) -> float: |
| | max_num_retries = 0 |
| | if max_out_len is None: |
| | max_out_len = 1 |
| | while max_num_retries < self.retry: |
| | self.wait() |
| | header = {'content-type': 'application/json'} |
| | try: |
| | input = self.input_format.replace('<input_text_to_replace>', |
| | input) |
| | data = dict(inputs=input, parameters=self.generation_kwargs) |
| | raw_response = requests.post(self.url, |
| | headers=header, |
| | data=json.dumps(data)) |
| | except requests.ConnectionError: |
| | self.logger.error('Got connection error, retrying...') |
| | continue |
| | try: |
| | response = raw_response.json() |
| |
|
| | assert ('prompt_token_ids' in response and 'prompt_logprobs' |
| | in response), f'prompt_token_ids and prompt_logprobs \ |
| | must be in the output. \ |
| | Please consider adding \ |
| | --return_all_prompt_logprobs argument \ |
| | when starting lightllm service. Response: {str(response)}' |
| |
|
| | prompt_token_ids = response['prompt_token_ids'][1:] |
| | prompt_logprobs = [ |
| | item[1] for item in response['prompt_logprobs'] |
| | ] |
| | logprobs = [ |
| | item[str(token_id)] for token_id, item in zip( |
| | prompt_token_ids, prompt_logprobs) |
| | ] |
| | if len(logprobs) == 0: |
| | return 0.0 |
| | ce_loss = -sum(logprobs) / len(logprobs) |
| | return ce_loss |
| | except requests.JSONDecodeError: |
| | self.logger.error('JsonDecode error, got', |
| | str(raw_response.content)) |
| | max_num_retries += 1 |
| | raise RuntimeError('Calling LightllmAPI failed after retrying for ' |
| | f'{max_num_retries} times. Check the logs for ' |
| | 'details.') |
| |
|
