AndresR2909 commited on
Commit
c57c722
1 Parent(s): 2c53521

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +38 -0
handler.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import pipeline
3
+
4
+ class EndpointHandler():
5
+ def __init__(self, path=""):
6
+ self.pipeline = pipeline(model=path, truncation=True,)
7
+
8
+
9
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
10
+ """
11
+ data args:
12
+ inputs (:obj: `str`)
13
+ date (:obj: `str`)
14
+ Return:
15
+ A :obj:`list` | `dict`: will be serialized and returned
16
+ """
17
+ # get inputs
18
+ input = data.get("inputs",data)
19
+ date = data.get("date", None)
20
+
21
+ # Realizar la inferencia
22
+ outputs = self.pipeline([input])
23
+
24
+ # Ajustar las etiquetas a 0 y 1
25
+ label_mapping = {"LABEL_0": 0, "LABEL_1": 1}
26
+ label_names = {0: "sin_intencion", 1: "intencion_suicida"}
27
+
28
+ # Modificar las salidas con el formato deseado
29
+ adjusted_results = [
30
+ {
31
+ "input": text,
32
+ "clasiffication": str(label_mapping[result['label']]), # "clasiffication" como string
33
+ "label": label_names[label_mapping[result['label']]] # Mapeo de etiquetas a nombres
34
+ }
35
+ for result, text in zip(results, texts)
36
+ ]
37
+
38
+ return adjusted_results