wasimmadha/entity-extraction

1import itertools2import torch 3import numpy as np4from tqdm.auto import tqdm5 6def get_char_probs(texts, predictions, tokenizer):7    """8    Maps prediction from encoded offset mapping to the text9 10    Prediction = 466 sequence length * batch11    text = 768 * batch12    Using offset mapping [(0, 4), ] -- 46613 14    creates results that is size of texts15 16    for each text result[i]17    result[0, 4] = pred[0] like wise for all18    19    """20    results = [np.zeros(len(t)) for t in texts]21    for i, (text, prediction) in enumerate(zip(texts, predictions)):22        encoded = tokenizer(text, 23                            add_special_tokens=True,24                            return_offsets_mapping=True)25        for idx, (offset_mapping, pred) in enumerate(zip(encoded['offset_mapping'], prediction)):26            start = offset_mapping[0]27            end = offset_mapping[1]28            results[i][start:end] = pred29    return results30 31 32def get_results(char_probs, th=0.5):33    """34      Get the list of probabilites with size of text35      And then get the index of the characters which are more than th36      example:37          char_prob = [0.1, 0.1, 0.9, 0.9, 0.9, 0.9, 0.2, 0.2, 0.2, 0.7, 0.7, 0.7] ## length == 76638          where > 0.5 index ## [ 2,  3,  4,  5,  9, 10, 11]39 40          Groupby same one -- [[2, 3, 4, 5], [9, 10, 11]]41          And get the max and min and output the results42 43    """44    results = []45    for char_prob in char_probs:46        result = np.where(char_prob >= th)[0] + 147        result = [list(g) for _, g in itertools.groupby(result, key=lambda n, c=itertools.count(): n - next(c))]48        result = [f"{min(r)} {max(r)}" for r in result]49        result = ";".join(result)50        results.append(result)51    return results52 53 54def get_predictions(results):55    """56      Will get the location, as a string, just like location in the df57      results = ['2 5', '9 11']58 59      loop through, split it and save it as start and end and store it in array60    """61    predictions = []62    for result in results:63        prediction = []64        if result != "":65            for loc in [s.split() for s in result.split(';')]:66                start, end = int(loc[0]), int(loc[1])67                prediction.append([start, end])68        predictions.append(prediction)69    return predictions70 71def inference_fn(test_loader, model, device):72    preds = []73    model.eval()74    model.to(device)75    tk0 = tqdm(test_loader, total=len(test_loader))76    for inputs in tk0:77        for k, v in inputs.items():78            inputs[k] = v.to(device)79        with torch.no_grad():80            y_preds = model(inputs)81        preds.append(y_preds.sigmoid().numpy())82    predictions = np.concatenate(preds)83    return predictions84 85def get_text(context, indexes):86  if (indexes):87      if ';' in indexes:88        list_indexes = indexes.split(';')89        90        answer = ''91        for idx in list_indexes:92          start_index = int(idx.split(' ')[0])93          end_index = int(idx.split(' ')[1])94          answer += ' ' 95          answer += context[start_index:end_index]96        return answer97      else:98        start_index = int(indexes.split(' ')[0])99        end_index = int(indexes.split(' ')[1])100 101        return context[start_index:end_index]102  else:103    return 'Not found in this Context'104 105