Update handler.py
Browse files- handler.py +8 -2
handler.py
CHANGED
@@ -19,10 +19,16 @@ class EndpointHandler():
|
|
19 |
self.tokenizer = AutoTokenizer.from_pretrained(cfg['repo'])
|
20 |
|
21 |
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
22 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
23 |
max_new_tokens = 1024
|
24 |
if "max_new_tokens" in data:
|
25 |
-
max_new_tokens = data
|
26 |
try:
|
27 |
max_new_tokens = int(max_new_tokens)
|
28 |
except Exception as e:
|
|
|
19 |
self.tokenizer = AutoTokenizer.from_pretrained(cfg['repo'])
|
20 |
|
21 |
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
22 |
+
"""
|
23 |
+
data args:
|
24 |
+
inputs (:obj: `str` | `PIL.Image` | `np.array`)
|
25 |
+
kwargs
|
26 |
+
Return:
|
27 |
+
A :obj:`list` | `dict`: will be serialized and returned
|
28 |
+
"""
|
29 |
max_new_tokens = 1024
|
30 |
if "max_new_tokens" in data:
|
31 |
+
max_new_tokens = data["max_new_tokens"]
|
32 |
try:
|
33 |
max_new_tokens = int(max_new_tokens)
|
34 |
except Exception as e:
|