Coverage for /var/devmt/py/docp-dbi_1.1.0/docp_dbi/databases/_embeddingfunctions.py: 100%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4:Purpose: This module provides a the custom embedding functions which 

5 enable a user to define the embedding function used for 

6 ChromaDB's embedding. 

7 

8 .. important:: 

9 

10 This module should *not* be interacted with directly. 

11 

12:Platform: Linux/Windows | Python 3.11+ 

13:Developer: J Berendt 

14:Email: development@s3dev.uk 

15 

16:Comments: n/a 

17 

18""" 

19# pylint: disable=wrong-import-order 

20 

21from __future__ import annotations 

22 

23import chromadb 

24import os 

25import torch 

26from sentence_transformers import SentenceTransformer 

27 

28# Module based constants: 

29MODEL_CACHE = os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))), 

30 '.cache') 

31 

32 

33class CustomEmbeddingFunction(chromadb.EmbeddingFunction): 

34 """Enable user to specify the embedding model to be used. 

35 

36 Args: 

37 embedding_model_path (str): Full path to the embedding model to 

38 be used. Note: This *must* be a valid repo path, as GGUF 

39 files are currently not supported by the 

40 ``SentenceTransformer`` class. 

41 local_files_only (bool): Whether or not to only look at local 

42 files (i.e., do not try to download the model). This is 

43 passed directly into the ``SentenceTransformer`` class. 

44 

45 

46 """ 

47 

48 # Note: Installing torch is a huge overhead, just for this. However, torch 

49 # will already be installed as part of the sentence-transformers library, 

50 # so we'll use it here too. 

51 _MODEL_KWARGS = {'device': 'cuda' if torch.cuda.is_available() else 'cpu', 

52 'trust_remote_code': True} 

53 

54 def __init__(self, embedding_model_path: str, local_files_only: bool): 

55 """Custom embedding function class initialiser.""" 

56 self._path = embedding_model_path 

57 self._model = SentenceTransformer(self._path, 

58 cache_folder=MODEL_CACHE, 

59 local_files_only=local_files_only, 

60 **self._MODEL_KWARGS) 

61 

62 def __call__(self, input: Document) -> Embeddings: # nocover # noqa # pylint: disable=undefined-variable 

63 """Required call method for a custom embedding function.""" 

64 # pylint: disable=redefined-builtin # Required by EmbeddingFunction 

65 return self.embed_documents(input=input) 

66 

67 @property 

68 def model(self): 

69 """Accessor to the internal embedding model.""" 

70 return self._model 

71 

72 @property 

73 def model_path(self) -> str: 

74 """Accessor to the requested model path.""" 

75 return self._path 

76 

77 def embed_documents(self, input: Document) -> Embeddings: # noqa # pylint: disable=undefined-variable 

78 """Create document embeddings. 

79 

80 This method is called by the :mod:`langchain_chroma.vectorstores` 

81 module when loading documents into the database. 

82 

83 """ 

84 # pylint: disable=redefined-builtin # Required by EmbeddingFunction 

85 return self._model.encode(sentences=input, device=self._MODEL_KWARGS.get('device'))