nizar-sayad commited on
Commit
3625af8
1 Parent(s): ef14c81

added custom handler

Browse files
Files changed (2) hide show
  1. handler.py +31 -0
  2. requirements.txt +3 -0
handler.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification
3
+ import torch
4
+
5
+
6
+
7
+ class EndpointHandler:
8
+ def __init__(self, path=""):
9
+ # load model and processor from path
10
+ guider_config = AutoConfig.from_pretrained(path)
11
+ self.model = AutoModelForSequenceClassification.from_pretrained(path, config=guider_config)
12
+ self.tokenizer = AutoTokenizer.from_pretrained(path)
13
+
14
+ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
15
+ """
16
+ Args:
17
+ data (:dict:):
18
+ The payload with the text prompt.
19
+ """
20
+ # process input
21
+ gen_outputs_no_input_decoded = data.pop("gen_outputs_no_input_decoded", data)
22
+
23
+ # Guiding the model with his ranking,
24
+ guider_inputs = self.tokenizer([gen_output_no_input_decoded for gen_output_no_input_decoded in gen_outputs_no_input_decoded],
25
+ return_tensors='pt', padding=True, truncation=True)
26
+
27
+ guider_outputs = self.model(**guider_inputs)
28
+
29
+ # the slicing at the end [:,x]: x=0 for negative, x=1 for neutral, x=2 for positive
30
+ guider_predictions = torch.nn.functional.softmax(guider_outputs.logits, dim=-1)[:, 0].tolist()
31
+ return {"guider_predictions": guider_predictions}
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ accelerate
2
+ bitsandbytes
3
+ transformers