File size: 760 Bytes
de7f614
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e11605b
 
de7f614
 
e11605b
de7f614
 
 
 
6c69c8a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
from typing import Dict, List, Any
from InstructorEmbedding import INSTRUCTOR

class EndpointHandler:
    def __init__(self, path=""):
        # load model on gpu
        self.model = INSTRUCTOR(path, device="cuda")

    def __call__(self, data: Dict[str, Any]) -> List[List[float]]:
        """
         data args:
              inputs (:obj: `str`)
        Return:
              A :obj:`list` | `dict`: will be serialized and returned
        """
        # get inputs
        inputs: dict = data.pop("inputs", data)

        print(inputs)

        # add instruction
        query  = [['Retrieve documents that can help answer the question:',
                   inputs]]

        # encode
        embedding = self.model.encode(query)

        return embedding