| | from typing import Dict, List, Any |
| | from transformers import AutoProcessor, MarkupLMModel |
| |
|
| |
|
| | class EndpointHandler(): |
| | def __init__(self, path=""): |
| | self.processor = AutoProcessor.from_pretrained("microsoft/markuplm-large") |
| | self.model = MarkupLMModel.from_pretrained("microsoft/markuplm-large") |
| |
|
| |
|
| | def __call__(self, data: Any) -> List[List[Dict[str, float]]]: |
| | """ |
| | Args: |
| | data (:obj:): |
| | includes the input data and the parameters for the inference. |
| | Return: |
| | A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing : |
| | - "label": A string representing what the label/class is. There can be multiple labels. |
| | - "score": A score between 0 and 1 describing how confident the model is for this label/class. |
| | """ |
| | |
| | inputs = data.pop("inputs", data) |
| | encoding = self.processor(inputs, return_tensors="pt") |
| | output = self.model(**encoding) |
| | return {"last_hidden_state": output.last_hidden_state[0].tolist(), |
| | "pooler_output": output.pooler_output[0].tolist()} |
| |
|
| |
|
