MichaelAI23
commited on
Commit
•
3aac9e2
1
Parent(s):
2a339d3
Update the handler file
Browse files- handler.py +2 -1
handler.py
CHANGED
@@ -1,5 +1,5 @@
|
|
1 |
from typing import Dict, Any
|
2 |
-
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
|
3 |
import torch
|
4 |
|
5 |
class EndpointHandler:
|
@@ -11,6 +11,7 @@ class EndpointHandler:
|
|
11 |
)
|
12 |
# load model and processor from path
|
13 |
self.model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config)
|
|
|
14 |
|
15 |
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
|
16 |
"""
|
|
|
1 |
from typing import Dict, Any
|
2 |
+
from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer
|
3 |
import torch
|
4 |
|
5 |
class EndpointHandler:
|
|
|
11 |
)
|
12 |
# load model and processor from path
|
13 |
self.model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config)
|
14 |
+
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
|
15 |
|
16 |
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
|
17 |
"""
|