ner.py

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']