Add example extractors

.gitignore

@@ -15,3 +15,4 @@ build
 upload.pdf
 indexify_extractor_sdk.egg-info
 extractors/indexify_extractors.egg-info/
+data/

extractors/identity_hash_embedding.py

@@ -0,0 +1,40 @@
+import hashlib
+import numpy as np
+from typing import List
+
+from indexify_extractor_sdk import (ExtractorSchema, EmbeddingSchema)
+from indexify_extractor_sdk.base_embedding import (BaseEmbeddingExtractor, EmbeddingInputParams)
+
+class IdentityHashEmbedding(BaseEmbeddingExtractor):
+    """
+        Implements a Hash Extractor, which can be used to find duplicates within the dataset.
+        It hashes the text into bytes, and interprets these are a numpy array.
+
+        We can extend this by LocalitySensitiveHashing, to also account for small perturbations in the input bytes.
+
+        This is equivalent to an identity mapping (with the sample-size n large enough, there will be collisions, but this is highly unlikely )
+    """
+
+    def __init__(self):
+        super(IdentityHashEmbedding, self).__init__(max_context_length=128)
+
+    def extract_embeddings(self, texts: List[str]) -> List[List[float]]:
+        return [self._embed(text) for text in texts]
+
+    def extract_query_embeddings(self, query: str) -> List[float]:
+        return self._embed(query)
+
+    def schemas(self) -> ExtractorSchema:
+        input_params = EmbeddingInputParams()
+        return ExtractorSchema(
+            input_params=input_params.model_dump_json(),
+            embedding_schemas={
+                "embedding": EmbeddingSchema(distance_metric="cosine", dim=32)
+            },
+        )
+
+    def _embed(self, text) -> List[float]:
+        model = hashlib.sha256()
+        model.update(bytes(text, 'utf-8'))
+        out = model.digest()
+        return np.frombuffer(out, dtype=np.int8).tolist()

extractors/identity_hash_embedding.yaml

@@ -0,0 +1,8 @@
+name: yenicelik/identity-hash-extractor
+version: 1
+description: "Embedding that extracts the Sha256 hash of whatever bytestring was inputted. This is useful to quickly check for duplicates."
+module: identity_hash_embedding:IdentityHashEmbedding
+gpu: false
+python_dependencies:
+  - numpy
+system_dependencies:

extractors/language_extractor.py

@@ -0,0 +1,56 @@
+import json
+from pydantic import BaseModel
+from typing import List, Literal
+
+from indexify_extractor_sdk import (
+    Extractor,
+    Feature,
+    ExtractorSchema,
+    Content,
+)
+from lingua import LanguageDetectorBuilder
+
+
+class LanguageExtractionInputParams(BaseModel):
+    overlap: int = 0
+    text_splitter: Literal["char", "token", "recursive", "new_line"] = "new_line"
+
+class LanguageExtractor(Extractor):
+    """
+    Extractor class for detecting the language of given content.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._model = LanguageDetectorBuilder.from_all_languages().build()
+
+    def extract(
+        self, content: List[Content], params: LanguageExtractionInputParams
+    ) -> List[List[Content]]:
+        content_texts = [c.data.decode("utf-8") for c in content]
+        out = []
+        for i, x in enumerate(content):
+            language = self._model.detect_language_of(content_texts[i])
+            confidence = self._model.compute_language_confidence(content_texts[i], language)
+            # TODO: Could be modified depending on the database we have
+            data = {"language": language.name, "score": str(confidence)}
+            out.append(
+                [Content.from_text(
+                    text=content_texts[i],
+                    feature=Feature.metadata(value=data, name="language"),
+                )]
+            )
+        return out
+
+    def schemas(self) -> ExtractorSchema:
+        """
+        Returns a list of options for indexing.
+        """
+        input_params = LanguageExtractionInputParams()
+        # TODO If it's metadata, how do we extract things
+        # This extractor does not return any embedding, only a dictionary!
+        return ExtractorSchema(
+            embedding_schemas={},
+            input_params=json.dumps(input_params.model_json_schema()),
+        )
+

extractors/language_extractor.yaml

@@ -0,0 +1,8 @@
+name: diptanu/language-extractor
+version: 1
+description: "Classifies text into a human language"
+module: language_extractor:LanguageExtractor
+gpu: false
+python_dependencies:
+  - lingua-language-detector
+system_dependencies:

extractors/pdf_embedder.py

@@ -32,7 +32,7 @@ def extract(
                 for (i, (chunk, embeddings)) in enumerate(zip(chunks, embeddings)):
                     embeddings_list.append(Embeddings(content_id=c.id, text=chunk, embeddings=embeddings, metadata=json.dumps({"page": i})))
         return embeddings_list
-    
+
     def extract_query_embeddings(self, query: str) -> List[float]:
         return self._model.embed_query(query)
 
@@ -44,4 +44,4 @@ def info(self) -> ExtractorInfo:
             input_params=input_params,
             output_schema=EmbeddingSchema(distance="cosine", dim=384),
         )
- 
+

extractors/simple_invoice_parser.py

@@ -0,0 +1,95 @@
+from io import BytesIO
+import json
+import re
+import requests
+import timeit
+from typing import List, Literal
+
+from PIL import Image
+from pdf2image import convert_from_bytes
+from pydantic import BaseModel
+import pytesseract
+import torch
+from transformers import DonutProcessor, VisionEncoderDecoderModel
+
+from indexify_extractor_sdk import (
+    Extractor,
+    Feature,
+    ExtractorSchema,
+    Content,
+)
+
+
+class SimpleInvoiceParserInputParams(BaseModel):
+    # No input except the file itself
+    ...
+
+class SimpleInvoiceParserExtractor(Extractor):
+    def __init__(self):
+        super().__init__()
+        self.processor = DonutProcessor.from_pretrained("to-be/donut-base-finetuned-invoices")
+        self.model = VisionEncoderDecoderModel.from_pretrained("to-be/donut-base-finetuned-invoices")
+        # TODO: Is this for example how we would pick it up? Probably the model would still need to be defined by the user i.e. how it should be used
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+
+    def _process_document(self, image):
+        # prepare encoder inputs
+        pixel_values = self.processor(image, return_tensors="pt").pixel_values
+
+        # prepare decoder inputs
+        task_prompt = "<s_cord-v2>"
+        decoder_input_ids = self.processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
+
+        # generate answer
+        outputs = self.model.generate(
+            pixel_values.to(self.device),
+            decoder_input_ids=decoder_input_ids.to(self.device),
+            max_length=self.model.decoder.config.max_position_embeddings,
+            early_stopping=True,
+            pad_token_id=self.processor.tokenizer.pad_token_id,
+            eos_token_id=self.processor.tokenizer.eos_token_id,
+            use_cache=True,
+            num_beams=1,
+            bad_words_ids=[[self.processor.tokenizer.unk_token_id]],
+            return_dict_in_generate=True,
+        )
+
+        # postprocess
+        sequence = self.processor.batch_decode(outputs.sequences)[0]
+        sequence = sequence.replace(self.processor.tokenizer.eos_token, "").replace(self.processor.tokenizer.pad_token, "")
+        sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
+        return self.processor.token2json(sequence), image
+
+    def extract(
+        self, content: List[Content], params: SimpleInvoiceParserInputParams
+    ) -> List[List[Content]]:
+        content_filebytes = [c.data for c in content]
+
+        # TODO: Right now it only looks at the first image! We should probably flatten it and do it for each page!
+        images = [convert_from_bytes(x)[0].convert("RGB") for x in content_filebytes]
+
+        out = []
+        for i, x in enumerate(content):
+            print("i, x are: ", i, x)
+            data = self._process_document(images[i])[0]  # Key 1 includes the image, which we ignore in this case
+            out.append(
+                [Content.from_text(
+                    text="",  # TODO: Diptanu, what do we do for PDFs? Do you want to save the raw bytes too, I feel like this is unnecessary? Also, I felt like these would be stored in a database _before_ processing, not after
+                    feature=Feature.metadata(value=data, name="invoice_simple_donut"),
+                )]
+            )
+        return out
+
+    def schemas(self) -> ExtractorSchema:
+        """
+        Returns a list of options for indexing.
+        """
+        input_params = SimpleInvoiceParserExtractor()
+        # TODO If it's metadata, how do we extract things
+        # This extractor does not return any embedding, only a dictionary!
+        return ExtractorSchema(
+            embedding_schemas={},
+            input_params=json.dumps(input_params.model_json_schema()),
+        )
+

extractors/simple_invoice_parser.yaml

@@ -0,0 +1,19 @@
+name: yenicelik/simple-invoice-parser
+version: 1
+description: "Parses an invoice using the to-be/donut-base-finetuned-invoices huggingface model"
+module: simple_invoice_parser:SimpleInvoiceParserExtractor
+gpu: true
+python_dependencies:
+  - transformers
+  - torch
+  - torchvision
+  - Pillow
+  - pytesseract
+  - timm
+  - sentencepiece
+  - donut-python
+  - pdf2image
+system_dependencies:
+  - tesseract-ocr
+  - libtesseract-dev
+  - poppler-utils