| from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer |
| from typing import Dict, List, Any |
| from tokenizers.decoders import WordPiece |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path="."): |
| model = AutoModelForTokenClassification.from_pretrained(path) |
| tokenizer = AutoTokenizer.from_pretrained(path) |
| self.pipeline = pipeline('ner', model=model, tokenizer=tokenizer, aggregation_strategy='simple') |
| self.pipeline.tokenizer.backend_tokenizer.decoder = WordPiece() |
|
|
| 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 |
| """ |
| return self.pipeline(data['inputs']) |
|
|
