|
from typing import Any, Dict, List |
|
|
|
from span_marker import SpanMarkerModel |
|
|
|
|
|
class EndpointHandler: |
|
def __init__(self, model_id: str) -> None: |
|
self.model = SpanMarkerModel.from_pretrained(model_id) |
|
|
|
self.model.try_cuda() |
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
""" |
|
Args: |
|
data (Dict[str, Any]): |
|
a dictionary with the "inputs" key corresponding to a string containing some text |
|
Return: |
|
A List[Dict[str, Any]]:. The object returned should be like [{"entity_group": "XXX", "word": "some word", "start": 3, "end": 6, "score": 0.82}] containing : |
|
- "entity_group": A string representing what the entity is. |
|
- "word": A rubstring of the original string that was detected as an entity. |
|
- "start": the offset within `input` leading to `answer`. context[start:stop] == word |
|
- "end": the ending offset within `input` leading to `answer`. context[start:stop] === word |
|
- "score": A score between 0 and 1 describing how confident the model is for this entity. |
|
""" |
|
return [ |
|
{ |
|
"entity_group": entity["label"], |
|
"word": entity["span"], |
|
"start": entity["char_start_index"], |
|
"end": entity["char_end_index"], |
|
"score": entity["score"], |
|
} |
|
for entity in self.model.predict(data["inputs"]) |
