1  """
 2  Hugging Face module
 3  """
 4  
 5  from ...models import Models, PoolingFactory
 6  
 7  from ..base import Vectors
 8  
 9  
10  class HFVectors(Vectors):
11      """
12      Builds vectors using the Hugging Face transformers library.
13      """
14  
15      @staticmethod
16      def ismethod(method):
17          """
18          Checks if this method uses local transformers-based models.
19  
20          Args:
21              method: input method
22  
23          Returns:
24              True if this is a local transformers-based model, False otherwise
25          """
26  
27          return method in ("transformers", "pooling", "clspooling", "meanpooling")
28  
29      def loadmodel(self, path):
30          # Build embeddings with transformers pooling
31          return PoolingFactory.create(
32              {
33                  "method": self.config.get("method"),
34                  "path": path,
35                  "device": Models.deviceid(self.config.get("gpu", True)),
36                  "tokenizer": self.config.get("tokenizer"),
37                  "maxlength": self.config.get("maxlength"),
38                  "loadprompts": "instructions" not in self.config,
39                  "modelargs": self.config.get("vectors", {}),
40              }
41          )
42  
43      def encode(self, data, category=None):
44          # Encode data using vectors model
45          return self.model.encode(data, batch=self.encodebatch, category=category)