| import torch
|
|
|
| from ScanDL2 import ScanDL2
|
|
|
|
|
| class EndpointHandler:
|
| def __init__(self, path: str = ""):
|
|
|
| self.models = {
|
| "sentence": ScanDL2(
|
| text_type="sentence",
|
| bsz=2,
|
| save=None,
|
| filename=None,
|
| ),
|
| "paragraph": ScanDL2(
|
| text_type="paragraph",
|
| bsz=2,
|
| save=None,
|
| filename=None,
|
| ),
|
| }
|
|
|
| for m in self.models.values():
|
|
|
| m.eval()
|
|
|
| def __call__(self, data):
|
|
|
| inputs = data.get("inputs", data)
|
|
|
| parameters = data.get("parameters", {})
|
|
|
| text_type = parameters.get("text_type", "sentence")
|
| model = self.models[text_type]
|
| bsz = parameters.get("bsz", 2)
|
|
|
| if model.scandl_module.args.batch_size != bsz:
|
| model.scandl_module.args.batch_size = bsz
|
| model.fixdur_module.bsz = bsz
|
| model.fixdur_module.args["bsz"] = bsz
|
|
|
| if isinstance(inputs, str):
|
| texts = [inputs]
|
| elif isinstance(inputs, list):
|
| texts = inputs
|
| else:
|
| raise ValueError("'inputs' must be a string or list of strings.")
|
|
|
| with torch.no_grad():
|
| output = model(texts=texts)
|
|
|
| return output
|
|
|
