| from simpletransformers.classification import ClassificationModel, ClassificationArgs | |
| from typing import Dict, List, Any | |
| import pandas as pd | |
| import webvtt | |
| from datetime import datetime | |
| import torch | |
| import spacy | |
| nlp = spacy.load("en_core_web_sm") | |
| tokenizer = nlp.tokenizer | |
| token_limit = 200 | |
| class EndpointHandler(): | |
| def __init__(self, path="."): | |
| print("Loading models...") | |
| cuda_available = torch.cuda.is_available() | |
| self.model = ClassificationModel( | |
| "roberta", path, use_cuda=cuda_available | |
| ) | |
| def __call__(self, data_file: str) -> List[Dict[str, Any]]: | |
| ''' data_file is a str pointing to filename of type .vtt ''' | |
| utterances_list = [] | |
| predictions, raw_outputs = self.model.predict(utterances_list) | |
| return predictions | |