File size: 614 Bytes
a8d332e a2b60e7 a8d332e 1690680 a8d332e a2b60e7 2a6910f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
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
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 '''
return []
|