llm-app

1
from llm_app.model_wrappers.base import BaseModel
2

3

4
class SentenceTransformerTask(BaseModel):
5
    def __init__(self, model: str, device: str = "cpu", **kwargs):
6
        """
7
        Wrapper for sentence-transformers.
8
        Arguments:
9
            model: model name or path
10
        """
11
        from sentence_transformers import SentenceTransformer
12

13
        super().__init__(**kwargs)
14
        self.model = SentenceTransformer(model_name_or_path=model, device=device)
15

16
    def __call__(self, text: str, **kwargs) -> str:
17
        """
18
        Arguments:
19
            text: input text string.
20
            **kwargs: Check out https://www.sbert.net/
21
        """
22
        return self.model.encode(text, **kwargs).tolist()
23