File size: 2,531 Bytes
5a85bf7 915446d 5a85bf7 316f722 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
from typing import Dict, List, Any
from setfit import SetFitModel
import numpy as np
class EndpointHandler:
def __init__(self, path=""):
# load model
self.model = SetFitModel.from_pretrained(path)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
data args:
inputs (:obj: `str` | `PIL.Image` | `np.array`)
kwargs
Return:
A :obj:`list` | `dict`: will be serialized and returned
"""
# get inputs
inputs = data.pop("inputs", data)
if isinstance(inputs, str):
inputs = [inputs]
exerciselabels = ['positive experience',
'power posing',
'worry vs rumination',
'self-confidence',
'negative emotions',
'sleep',
'loneliness',
'imaginary friend',
'perfectionism',
'negative self-talk',
'woop',
'venting',
'worry window',
'act of kindness',
'blowing balloons',
'feeling on anger',
'power of smile',
'body scan',
'stress enhancing thoughts',
'anger ball of fire',
'emotions',
'lean against wall',
'breathing',
'crossed arms',
'energy traffic light',
'boundaries',
'Inner strength']
# run normal prediction
preds = self.model.predict(inputs)
scores = self.model.predict_proba(inputs)
label = [[el for el, p in zip(exerciselabels, ps) if p] for ps in preds]
# Modify the label array
modified_label = label[0]
# Extract the positives probabilities from each inner array
modified_proba = [[inner[0][1]] for item, inner in zip(scores, scores)]
score = [[el for el, p in zip(modified_proba, ps) if p] for ps in preds]
# Modify the score array
modified_score = score[0]
combined_dict = {key: value for key, value in zip(modified_label, modified_score)}
output_array = [combined_dict]
# for element in combined_array:
return label |