Ben Nguyen
commited on
Commit
•
1320603
1
Parent(s):
3146dc1
- handler.py +3 -5
handler.py
CHANGED
@@ -1,13 +1,11 @@
|
|
1 |
from typing import Dict, List, Any
|
2 |
-
from comet import load_from_checkpoint
|
3 |
import os
|
4 |
|
5 |
|
6 |
class EndpointHandler():
|
7 |
def __init__(self, path=""):
|
8 |
-
|
9 |
-
|
10 |
-
self.model = load_from_checkpoint("model.ckpt")
|
11 |
|
12 |
def __call__(self, data: Dict[str, Any]) -> List[Any]:
|
13 |
"""
|
@@ -24,7 +22,7 @@ class EndpointHandler():
|
|
24 |
|
25 |
data = inputs.pop("data")
|
26 |
|
27 |
-
model_output = self.
|
28 |
scores = model_output["scores"]
|
29 |
|
30 |
return scores
|
|
|
1 |
from typing import Dict, List, Any
|
2 |
+
from comet import load_from_checkpoint as lfc
|
3 |
import os
|
4 |
|
5 |
|
6 |
class EndpointHandler():
|
7 |
def __init__(self, path=""):
|
8 |
+
self.m = lfc("model.ckpt")
|
|
|
|
|
9 |
|
10 |
def __call__(self, data: Dict[str, Any]) -> List[Any]:
|
11 |
"""
|
|
|
22 |
|
23 |
data = inputs.pop("data")
|
24 |
|
25 |
+
model_output = self.m.predict(data, batch_size=batch_size, num_workers=workers, gpus=0)
|
26 |
scores = model_output["scores"]
|
27 |
|
28 |
return scores
|