Update handler.py
Browse files- handler.py +2 -0
handler.py
CHANGED
@@ -1,6 +1,7 @@
|
|
1 |
import os
|
2 |
import torch
|
3 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
|
|
4 |
|
5 |
class EndpointHandler():
|
6 |
def __init__(self, model_id: str):
|
@@ -48,6 +49,7 @@ class EndpointHandler():
|
|
48 |
print(f"Configured label order for output: {self.label_names}")
|
49 |
|
50 |
|
|
|
51 |
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any] | List[Dict[str, Any]]:
|
52 |
"""
|
53 |
Handles inference requests.
|
|
|
1 |
import os
|
2 |
import torch
|
3 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
4 |
+
from typing import Dict, List, Any # <-- ADD THIS LINE
|
5 |
|
6 |
class EndpointHandler():
|
7 |
def __init__(self, model_id: str):
|
|
|
49 |
print(f"Configured label order for output: {self.label_names}")
|
50 |
|
51 |
|
52 |
+
# Corrected type hints in the signature below
|
53 |
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any] | List[Dict[str, Any]]:
|
54 |
"""
|
55 |
Handles inference requests.
|