import requests import json import os def get_completion(inputs, parameters=None): headers = { "Authorization": f"Bearer {model_api}", "Content-Type": "application/json" } data = { "inputs": inputs } if parameters is not None: data.update({"parameters": parameters}) response = requests.request("POST", model_url, headers=headers, data=json.dumps(data) ) return json.loads(response.content.decode("utf-8")) def merge_tokens(tokens): merged_tokens=[] for token in tokens: if merged_tokens and merged_tokens[-1]['entity_group'] == token['entity_group']: last_token = merged_tokens[-1] last_token['word'] += token['word'].replace('##', '') last_token['end'] = token['end'] last_token['score'] = (last_token['score'] + token['score']) / 2 else: merged_tokens.append(token) return merged_tokens def ner(input): output = get_completion(input) output = merge_tokens(output) return {"text": input, "entities": output} model_url = os.environ['API_URL'] model_api = os.environ['API_Key']