rara-tools 0.0.10__py3-none-any.whl → 0.0.11__py3-none-any.whl

rara_tools/digar_schema_converter.py

@@ -3,7 +3,7 @@ from typing import List, NoReturn
 
 from rara_tools.utils import lang_to_iso639_2, ratio_to_percentage
 
-GENERAL_DOC_IDENTIFIER = "Filepath"
+
 UNDEFINED_LANGUAGE_VALUE = "unk"
 QUALITY_RATIO_TYPE = "Float"
 
@@ -104,13 +104,15 @@ class DocSchemas:
             generated_id: str = "",
             permalink: str = "",
             min_language_ratio: float = 0.2,
-            convert_ratio: bool = True
+            convert_ratio: bool = True,
+            generated_id_type: str = "CustomID"
     ) -> NoReturn:
         self.__convert_ratio = convert_ratio
         self.__min_language_ratio = min_language_ratio
         self.__sierra_id = sierra_id
         self.__generated_id = generated_id
         self.__permalink = permalink
+        self.__generated_id_type = generated_id_type
         self.__doc_meta = doc_meta
         self.__ocr_accuracy_schema: dict = {}
         self.__text_quality_schema: dict = {}
@@ -209,7 +211,7 @@ class DocSchemas:
                 identifiers.append(
                     {
                         "@type": "Identifier",
-                        "qualifier": GENERAL_DOC_IDENTIFIER,
+                        "qualifier": self.__generated_id_type,
                         "value": self.__generated_id
                     }
                 )
@@ -235,6 +237,7 @@ class DIGARSchemaConverter:
             generated_id: str,
             sierra_id: str = "",
             permalink: str = "",
+            generated_id_type: str = "CustomID",
             min_language_ratio: float = 0.2,
             convert_ratio: bool = False
     ) -> NoReturn:
@@ -250,6 +253,8 @@ class DIGARSchemaConverter:
             Document's corresponding Sierra ID.
         permalink: str
             Permanent link, where the document can be accessed.
+        generated_id_type: str
+            Method / type of generated ID (e.g. 'UUID')
         min_language_ratio: float
             Cutoff ratio for languages. If ratio for some language
             does not exceed the set threshold, the language will not
@@ -264,6 +269,7 @@ class DIGARSchemaConverter:
         self.__sierra_id: str = sierra_id
         self.__generated_id: str = generated_id
         self.__permalink: str = permalink.removesuffix("/")
+        self.__generated_id_type: str = generated_id_type
         self.__texts: List[dict] = []
         self.__images: List[dict] = []
         self.__doc_meta: dict = {}
@@ -281,7 +287,8 @@ class DIGARSchemaConverter:
             generated_id=self.__generated_id,
             permalink=self.__permalink,
             min_language_ratio=self.__min_language_ratio,
-            convert_ratio=self.__convert_ratio
+            convert_ratio=self.__convert_ratio,
+            generated_id_type=self.__generated_id_type
         )
         self.__digar_schema: dict = {}
 

rara_tools/elastic.py

@@ -1,10 +1,11 @@
-from typing import Any, Dict, Iterator, Optional
+from typing import Any, Dict, Iterator, Optional, List
 
 import elasticsearch_dsl
 from elastic_transport import ObjectApiResponse
 from elasticsearch import Elasticsearch
 from elasticsearch.helpers import bulk
 from elasticsearch_dsl import Index
+from elasticsearch_dsl.response import Response
 
 from .decorators import _elastic_connection
 
@@ -82,6 +83,63 @@ class KataElastic:
     def add_mapping(self, index_name: str, schema: dict):
         index = Index(name=index_name)
         return index.put_mapping(body=schema, using=self.elasticsearch)
+    
+    
+    @_elastic_connection
+    def add_vector_mapping(
+            self, 
+            index_name: str, 
+            field: str, 
+            schema: Optional[dict] = None, 
+            dims: int = 1024
+    ) -> dict:
+        vector_mapping = {
+            "properties": {
+                field: {
+                    "type": "dense_vector", 
+                    "dims": dims
+                }
+            }
+        }
+        mapping = schema or vector_mapping
+        index = Index(name=index_name)
+        return index.put_mapping(body=mapping, using=self.elasticsearch)
+    
+    
+    @_elastic_connection
+    def add_ann_vector_mapping(
+            self, 
+            index_name: str,
+            field: str,
+            schema: Optional[dict] = None, 
+            dims: int = 1024
+    ) -> dict:
+        vector_mapping = {
+            "properties": {
+                field: {
+                    "type": "dense_vector", 
+                    "dims": dims, 
+                    "similarity": "cosine", 
+                    "index": True
+                }
+            }
+        }
+        mapping = schema or vector_mapping
+        index = Index(name=index_name)
+        return index.put_mapping(body=mapping, using=self.elasticsearch)
+
+    @_elastic_connection
+    def add_vector(
+            self, 
+            index_name: str,
+            document_id: str, 
+            vector: List[float], 
+            field: str
+    ) -> dict:
+        schema = {"doc": {field: vector}}
+        return self.elasticsearch.update(
+            index=index_name, id=document_id, body=schema, refresh="wait_for"
+        )
 
     @_elastic_connection
     def create_index(
@@ -170,6 +228,87 @@ class KataElastic:
             s.scan(), key=lambda doc: [getattr(doc, field) for field in sort_fields]
         )
         return documents
+    
+    
+    @_elastic_connection
+    def execute_fuzzy_search(
+            self, 
+            index: str,
+            field: str,
+            entity: str, 
+            fuzziness: int = 2, 
+            prefix_length: int = 1, 
+            max_expansions: int = 50
+    ) -> Response:
+        """Executes a fuzzy search.
+        :param: index str: Index to search from.
+        :param: entity str: Entity to search matches for.
+        :param: fuzziness int: Maximum edit distance for a match.
+        :param: prefix_length int: Number of characters in the prefix that 
+            should overlap with the original entity's prefix.
+        :param: max_expansion int: maximum number of terms the fuzzy query 
+            will match before halting the search
+        :return: Dict on search results.
+        """  
+        query_params = {
+            f"{field}.keyword": {
+                "value": entity,
+                "fuzziness": fuzziness,
+                "max_expansions": max_expansions,
+                "prefix_length": prefix_length
+            }
+        }
+        s = elasticsearch_dsl.Search(using=self.elasticsearch, index=index)
+        s = s.query("fuzzy", **query_params)
+        response = s.execute()
+        return response
+    
+    def execute_vector_search(
+            self,
+            index: str,
+            field: str,
+            query_vector: List[float],
+            k: int = 10,
+            num_candidates: int = 100,
+            n_docs: int = 10,
+            elastic_ids: List[str] = []
+    ) -> Response:
+        """ Execute a vector search.
+        NB! Works only with ANN mapping!
+        
+        :param: index str: Index to search from.
+        :param: field str: Field containing vectorized data.
+        :param: query vector List[float]: Vector to search matches for.
+        :param: k int: Number of nearest neighbors to return.
+        :param: num_candidates int: Number of candidates considered before selecting k results.
+        :param: elastic_ids: List[str]: Elastic ID-s for restricting the search.
+        """
+
+        s = elasticsearch_dsl.Search(using=self.elasticsearch, index=index)
+
+        # Add kNN vector search
+        s = s.extra(
+            knn={
+                "field": field,
+                "query_vector": query_vector, 
+                "k": k, 
+                "num_candidates": num_candidates
+            }
+        )
+
+        # Add ID filtering, if elastic_ids are specified
+        if elastic_ids:
+            s = s.query(
+                elasticsearch_dsl.Q("terms", _id=elastic_ids)
+            )
+            
+        # Sort by score and return `n_docs` best-matching documents
+        s = s.extra(size=n_docs)
+
+        # Execute the search
+        response = s.execute()
+        return response
+
 
     def __str__(self) -> str:
         return self.elasticsearch_url

{rara_tools-0.0.10.dist-info → rara_tools-0.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: rara-tools
-Version: 0.0.10
+Version: 0.0.11
 Summary: Tools to support Kata's work.
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10

{rara_tools-0.0.10.dist-info → rara_tools-0.0.11.dist-info}/RECORD

@@ -1,7 +1,7 @@
 rara_tools/converters.py,sha256=JcS74VzV6jm12l3C6aqMJBY9nuVW_aevQeCe32KmfrE,1576
 rara_tools/decorators.py,sha256=MjOyvZ5nTkwxwx2JLFEGpKKBysvecFw6EN6UDrSvZLU,2187
-rara_tools/digar_schema_converter.py,sha256=gGwhqdwxyTXODF0LP5Xi0u8uRoICfaIU3MRe1EVBnEc,13935
-rara_tools/elastic.py,sha256=vEvrbIPRtdqTdrNrPH2cewHLMfOTSf87a4JOiRQgYyA,7146
+rara_tools/digar_schema_converter.py,sha256=k95U2iRlEA3sh772-v6snhHW6fju6qSTMnvWJ6DpzZk,14254
+rara_tools/elastic.py,sha256=LZfHZqeTDjCEb5YX4CLPJEFffRSZAcRq6AtyP49Fo0E,11575
 rara_tools/exceptions.py,sha256=BwNh4qWxau_ylr9RqZoYwd1KnExI6oWWWDno3jkh8q4,474
 rara_tools/s3.py,sha256=uNDu2HzMYHAWh33RcHeyPFK7gdQfQPxsdfohyIKezEY,4467
 rara_tools/task_reporter.py,sha256=WCcZts9dAUokPc4vbrG3-lNAFLnWaMgE3b3iaUB7mr8,3256
@@ -9,8 +9,8 @@ rara_tools/utils.py,sha256=9vSbmuWYU5ydr4lXBKlUKa0xzDccFsaJv4T-XwgUfuY,2578
 rara_tools/constants/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rara_tools/constants/digitizer.py,sha256=gJ3jOMwuZfKcLqgOAxTyB266VYsskLabJiMUiSz3xX4,297
 rara_tools/constants/general.py,sha256=E9Jaw-YxocS_tOZw9QBoxO3e9KK5EMbLoM0R7D4Iflw,171
-rara_tools-0.0.10.dist-info/LICENSE.md,sha256=hkZVnIZll7e_KNEQzeY94Y9tlzVL8iVZBTMBvDykksU,35142
-rara_tools-0.0.10.dist-info/METADATA,sha256=jV6nZKhjjwDL6TWt-fKWudWNUAViZTVDL0J39fefFtM,3895
-rara_tools-0.0.10.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-rara_tools-0.0.10.dist-info/top_level.txt,sha256=JwfB5b8BAtW5OFKRln2AQ_WElTRyIBM4nO0FKN1cupY,11
-rara_tools-0.0.10.dist-info/RECORD,,
+rara_tools-0.0.11.dist-info/LICENSE.md,sha256=hkZVnIZll7e_KNEQzeY94Y9tlzVL8iVZBTMBvDykksU,35142
+rara_tools-0.0.11.dist-info/METADATA,sha256=pDcladCQ1A9O9Wh4UDSh0eHwNyqcGY1BOwPxSJKpLFk,3895
+rara_tools-0.0.11.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+rara_tools-0.0.11.dist-info/top_level.txt,sha256=JwfB5b8BAtW5OFKRln2AQ_WElTRyIBM4nO0FKN1cupY,11
+rara_tools-0.0.11.dist-info/RECORD,,