"""JSON-rpc proxy model for BIG-Bench.""" # Copyright 2021 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os from typing import Optional, Any, Dict, Union, List import urllib.parse import uuid import bigbench.api.model as model import requests import requests_unixsocket import time _RPC_VERSION = '2.0' _INET_URL_PREFIX = 'http://' _UNIX_URL_PREFIX = 'http+unix://' _MAX_RETRIES = 8 _RETRY_DELAY_SEC = 10 class JsonRpcModel(model.Model): """A BIGBench api-compatible json-rpc interface. Attributes: model_name: name of model url: json rpc url """ def __init__( self, model_name: str, host: Optional[str] = None, port: Optional[int] = None, socket_path: Optional[str] = None, ): """Initializes instance. Either host+port or socket_path must be specified. Args: model_name: name of model host: hostname for inet connections port: port number for inet connections socket_path: path to socket for unix-domain socket connections Raises: ValueError if host+port or socket_path is not set """ if socket_path is None: if host is None or port is None: raise ValueError('either host+port or socket_path must be specified') self.url = f'{_INET_URL_PREFIX}{host}:{port}/jsonrpc' else: abs_socket_path = os.path.abspath(os.path.expanduser(socket_path)) encoded_path = urllib.parse.quote_plus(abs_socket_path) self.url = f'{_UNIX_URL_PREFIX}{encoded_path}/jsonrpc' self.model_name = model_name def _post(self, payload: Dict[str, Any]) -> Optional[Any]: """Posts request to server. Args: payload: dictionary containing rpc params Returns: result from rpc call Raises: RuntimeError: An error from the server. HTTPError: An error communicating with server. """ for i in range(_MAX_RETRIES): if self.url.startswith(_UNIX_URL_PREFIX): with requests_unixsocket.monkeypatch(): response = requests.post(self.url, json=payload) else: response = requests.post(self.url, json=payload) if response.status_code != requests.codes.ok: #response.raise_for_status() print(f'error from server: {response}') time.sleep(_RETRY_DELAY_SEC) continue json_response = response.json() if 'result' not in json_response: #raise RuntimeError(f'response from server: {json_response}') print(f'error: bad response from server: {json_response}') time.sleep(_RETRY_DELAY_SEC) continue #print(f'response:\n{json_response["result"]}') return json_response['result'] raise RuntimeError(f'error: retry count exceeded') def _payload(self, method: str, params: Dict[str, Any]) -> Dict[str, Any]: """Generates payload dictionary for rpc call. Args: method: rpc method name params: parameter dictonary for rpc call Returns: json-serializable payload dictionary """ return { 'method': method, 'params': {**params, 'model_name': self.model_name}, 'jsonrpc': _RPC_VERSION, 'id': uuid.uuid1().int, } def generate_text( self, inputs: Union[str, List[str]], max_length: Optional[int] = None, stop_string: Optional[str] = None, output_regex: Optional[str] = None, ) -> Union[str, List[str]]: """Generates and returns outputs from language model. Args: inputs: String or list of input strings. max_length: Maximum output length, if None, limited only by model max output length stop_string: If specified, model output will be truncated to the shortest string which includes stop_string. output_regex: If specified, the first match to the python regular expression output_regex in the model output will be returned. If there is no match, an empty string will be returned. Returns: list of lists (of length=num_outputs) with generated responses Raises: RuntimeError: An error from the server. HTTPError: An error communicating with the server. """ return self._post( self._payload( method='generate_text', params={ 'inputs': inputs, 'max_length': max_length, 'stop_string': stop_string, 'output_regex': output_regex, })) def cond_log_prob( self, inputs: Union[str, List[str]], targets: Union[List[str], List[List[str]]], absolute_normalization: Optional[bool] = False, ) -> Union[List[float], List[List[float]]]: """Computes conditional log probabilities of targets given inputs. Args: inputs: A single string input or a list of string inputs. targets: Possible string outputs for each input. If input is a string, this is a list `[t_1, t_2, ..., t_n]` of possible string outputs. If input is a list of strings, then this is a nested list `[[t_1, t_2, ..., t_n], ...]` with length equal to `len(inputs)`. absolute_normalization: When True, the function returns the log probability of unconstrained generation or the target sequence. When False (default), log probabilities are normalized so that the probabilities of generating `targets` sum to 1. Note that setting `absolute_normalization` to True restricts the class of models that can be evaluated to those that can assign absolute probabilities to sequences. Returns: If a single string input is provided, returns a list of log-probabilities `[lp_1, lp_2, ..., lp_n]` predicted by the model, where `lp_i = log(prob(t_i | input)` is the conditional log-prob to generate target `t_i` given input. If a list of string inputs was provided, returns a list of such elements of the form `[[lp_1, lp_2, ..., lp_n], ...]`, where each element contains the log-probabilities for the corresponding input and targets. In this case, the length of the returned list is `len(input)`. If conditional probabilities are not supported by the model, the model returns None. Raises: RuntimeError: An error from the server. HTTPError: An error communicating with the server. """ return self._post( self._payload( method='cond_log_prob', params={ 'inputs': inputs, 'targets': targets, 'absolute_normalization': absolute_normalization, })) def shutdown_server(self) -> None: """Shuts down remote proxy server gracefully.""" self._post(self._payload(method='shutdown', params={})) def model_data(self) -> model.ModelData: """Returns ModelData for this model instance. Returns: ModelData. """ model_dict = self._post(self._payload(method='model_data', params={})) return model.ModelData(**model_dict)