Bashir Rastegarpanah commited on
Commit
78e1841
1 Parent(s): c710467

add custom handler

Browse files
Files changed (1) hide show
  1. handler.py +3 -5
handler.py CHANGED
@@ -1,16 +1,14 @@
1
  from typing import Dict, List, Any
2
- from transformers import AutoTokenizer
3
 
4
 
5
  class EndpointHandler:
6
  def __init__(self,
7
- model,
8
- # path=""
9
  ):
10
  # Preload all the elements you are going to need at inference.
11
  # pseudo:
12
- # self.model= load_model(path)
13
- self.model = model
14
  self.tokenizer = AutoTokenizer.from_pretrained("roberta-large", padding_side='right')
15
 
16
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
 
1
  from typing import Dict, List, Any
2
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
 
4
 
5
  class EndpointHandler:
6
  def __init__(self,
7
+ path="BashirRP/llm_judge2"
 
8
  ):
9
  # Preload all the elements you are going to need at inference.
10
  # pseudo:
11
+ self.model = AutoModelForSequenceClassification.from_pretrained(path)
 
12
  self.tokenizer = AutoTokenizer.from_pretrained("roberta-large", padding_side='right')
13
 
14
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: