cguynup commited on
Commit
2601994
·
1 Parent(s): 8256741

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +7 -2
handler.py CHANGED
@@ -1,8 +1,13 @@
1
- from transformers import pipeline
 
 
2
 
3
  class EndpointHandler():
4
  def __init__(self, path=""):
5
- self.pipeline = pipeline("text-classification", model=path)
 
 
 
6
 
7
 
8
  def __call__(self, data):
 
1
+ # Load model directly
2
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
3
+
4
 
5
  class EndpointHandler():
6
  def __init__(self, path=""):
7
+ tokenizer = AutoTokenizer.from_pretrained(path)
8
+ model = AutoModelForSequenceClassification.from_pretrained(path)
9
+
10
+ self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
11
 
12
 
13
  def __call__(self, data):