indexify 0.0.39__tar.gz → 0.0.40__tar.gz

{indexify-0.0.39 → indexify-0.0.40}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: indexify
-Version: 0.0.39
+Version: 0.0.40
 Summary: Python Client for Indexify
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: cloudpickle (>=3,<4)
 Requires-Dist: httpx[http2] (>=0,<1)
 Requires-Dist: pydantic (>=2.8,<3.0)
 Requires-Dist: pyyaml (>=6,<7)

{indexify-0.0.39 → indexify-0.0.40}/indexify/__init__.py

@@ -9,9 +9,11 @@ from .client import (
 from . import extractor_sdk
 from .settings import DEFAULT_SERVICE_URL
 from . import data_loaders
+from .graph import Graph
 
 __all__ = [
     "data_loaders",
+    "Graph",
     "Document",
     "extractor_sdk",
     "IndexifyClient",

{indexify-0.0.39 → indexify-0.0.40}/indexify/data_loaders/url_loader.py

@@ -1,7 +1,6 @@
 from . import DataLoader, FileMetadata
 from typing import List
 import httpx
-import hashlib
 import email.utils
 
 

{indexify-0.0.39 → indexify-0.0.40}/indexify/extractor_sdk/data.py

@@ -103,6 +103,6 @@ class ContentMetadata(BaseModel):
         )
 
 
-class PDFFile(BaseData):
+class File(BaseData):
     data: bytes
     mime_type: str

{indexify-0.0.39 → indexify-0.0.40}/indexify/extractor_sdk/extractor.py

@@ -10,7 +10,8 @@ import requests
 
 class EmbeddingSchema(BaseModel):
     dim: int
-    distance: str = "cosine"
+    distance: Optional[str] = "cosine"
+    database_url: Optional[str] = None
 
 class ExtractorMetadata(BaseModel):
     name: str
@@ -40,8 +41,8 @@ class Extractor(ABC):
 
     input_mime_types = ["text/plain"]
 
-    embeddings: Dict[str, EmbeddingSchema] = {}
-
+    embedding_indexes: Dict[str, EmbeddingSchema] = {}
+    
     @abstractmethod
     def extract(
         self, input: Type[BaseModel], params: Type[BaseModel] = None
@@ -55,31 +56,9 @@ class Extractor(ABC):
         pass
 
     @classmethod
-    @abstractmethod
     def sample_input(cls) -> Tuple[Content, Type[BaseModel]]:
         pass
 
-    def describe(self) -> ExtractorMetadata:
-        embedding_schemas = {}
-        try:
-            embedding_schemas = self.embedding_schemas
-        except NotImplementedError:
-            pass
-
-        json_schema = (
-            self._param_cls.model_json_schema() if self._param_cls is not None else None
-        )
-        return ExtractorMetadata(
-            name=self.name,
-            version=self.version,
-            description=self.description,
-            system_dependencies=self.system_dependencies,
-            python_dependencies=self.python_dependencies,
-            input_mime_types=self.input_mime_types,
-            embedding_schemas=embedding_schemas,
-            input_params=json.dumps(json_schema),
-        )
-
     def _download_file(self, url, filename):
         if os.path.exists(filename):
             # file exists skip
@@ -190,7 +169,7 @@ def extractor(
     python_dependencies: Optional[List[str]] = None,
     system_dependencies: Optional[List[str]] = None,
     input_mime_types: Optional[List[str]] = None,
-    embedding_schemas: Optional[Dict[str, EmbeddingSchema]] = None,
+    embedding_indexes: Optional[Dict[str, EmbeddingSchema]] = None,
     sample_content: Optional[Callable] = None,
 ):
     args = locals()
@@ -198,7 +177,7 @@ def extractor(
 
     def construct(fn):
         def wrapper():
-            hint = get_type_hints(fn).get("params", dict)
+            description = fn.__doc__ or args.get("description", "")
 
             if not args.get("name"):
                 args[
@@ -220,6 +199,7 @@ def extractor(
 
             for key, val in args.items():
                 setattr(DecoratedFn, key, val)
+            DecoratedFn.description = description
 
             return DecoratedFn
 

indexify-0.0.40/indexify/extractors/embedding.py

@@ -0,0 +1,53 @@
+from typing import List
+
+from indexify.extractor_sdk.data import Feature
+import torch
+import torch.nn.functional as F
+from transformers import AutoModel, AutoTokenizer
+from indexify.extractor_sdk.extractor import Extractor , Feature
+
+class SentenceTransformersEmbedding:
+    def __init__(self, model_name) -> None:
+        self._model_name = model_name
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            f"sentence-transformers/{model_name}"
+        )
+        self._model = AutoModel.from_pretrained(
+            f"sentence-transformers/{model_name}", torchscript=True
+        )
+        self._model.eval()
+
+    def embed_batch(self, inputs: List[str]) -> List[List[float]]:
+        result = self._embed(inputs)
+        return result.tolist()
+
+    def embed(self, query: str) -> List[float]:
+        result = self._embed([query])
+        return result[0].tolist()
+
+    def _embed(self, inputs: List[str]) -> torch.Tensor:
+        encoded_input = self._tokenizer(
+            inputs, padding=True, truncation=True, return_tensors="pt"
+        )
+        sentence_embeddings = self._model(**encoded_input)
+        return F.normalize(sentence_embeddings, p=2, dim=1)
+    
+class BasicSentenceTransformerModels(Extractor):
+
+    def __init__(self, model: str):
+        super().__init__()
+        self.model = SentenceTransformersEmbedding(model)
+
+    def extract(self, input: str) -> List[Feature]:
+        embeddings = self.model.embed(input)
+        return [Feature.embedding(values=embeddings)]
+    
+class BasicHFTransformerEmbeddingModels(Extractor):
+    
+        def __init__(self, model: str):
+            super().__init__()
+            self._model = AutoModel.from_pretrained(model, trust_remote_code=True)
+    
+        def extract(self, input: str) -> List[Feature]:
+            embeddings = self.model.embed_query(input)
+            return [Feature.embedding(values=embeddings)]

indexify-0.0.39/indexify/run_graph.py → indexify-0.0.40/indexify/graph.py

@@ -1,13 +1,12 @@
 import json
-
-from .extractor_sdk import Content, extractor, Extractor
-
+import itertools
 from collections import defaultdict
 from typing import Any, Dict, List, Optional, Type, Union
-from pydantic import BaseModel
 
-import itertools
+import cloudpickle
+from pydantic import BaseModel
 
+from .extractor_sdk import Content, extractor, Extractor
 from .runner import Runner
 
 @extractor(description="id function")
@@ -15,7 +14,10 @@ def _id(content: Content) -> List[Content]:
     return [content]
 
 
-class RunGraph:
+def load_graph(graph: bytes) -> 'Graph':
+    return cloudpickle.loads(graph)
+
+class Graph:
     def __init__(self, name: str, input: Type[BaseModel], start_node: extractor, runner: Runner):
         # TODO check for cycles
         self.name = name
@@ -35,7 +37,10 @@ class RunGraph:
 
         self.runner = runner
 
-    def _node(self, extractor: Union[extractor, Extractor], params: Any = None) -> 'RunGraph':
+    def get_extractor(self, name: str) -> Extractor:
+        return self.nodes[name]
+
+    def _node(self, extractor: Extractor, params: Any = None) -> 'Graph':
         name = extractor.name
 
         # if you've already inserted a node just ignore the new insertion.
@@ -49,13 +54,16 @@ class RunGraph:
         self._topo_counter[name] = 1
 
         return self
+    
+    def serialize(self):
+        return cloudpickle.dumps(self)
 
     def add_edge(
         self,
-        from_node: extractor,
-        to_node: extractor,
+        from_node: Type[Extractor],
+        to_node: Type[Extractor],
         prefilter_predicates: Optional[str] = None,
-    ) -> 'RunGraph':
+    ) -> 'Graph':
 
         self._node(from_node)
         self._node(to_node)
@@ -79,7 +87,7 @@ class RunGraph:
         from_node: extractor,
         to_nodes: List[extractor],
         prefilter_predicates: List[str] = [],
-    ) -> 'RunGraph':
+    ) -> 'Graph':
         print(f"{to_nodes}, {prefilter_predicates}, {prefilter_predicates}")
         for t_n, p in itertools.zip_longest(
             to_nodes, prefilter_predicates, fillvalue=None
@@ -99,7 +107,6 @@ class RunGraph:
 
     def run(self, wf_input, local):
         self._assign_start_node()
-        # self.runner = LocalRunner()
         self.runner.run(self, wf_input=wf_input)
         pass
 

{indexify-0.0.39 → indexify-0.0.40}/indexify/local_runner.py

@@ -10,7 +10,7 @@ from indexify.extractor_sdk.extractor import extractor, Extractor
 from collections import defaultdict
 from typing import Any, Callable, Dict, Optional, Union
 
-from indexify.run_graph import RunGraph
+from indexify.graph import Graph
 from indexify.runner import Runner
 
 
@@ -27,7 +27,7 @@ class LocalRunner(Runner):
     # those bytes have to be a python type
 
     # _input needs to be serializable into python object (ie json for ex) and Feature
-    def _run(self, g: RunGraph, _input: BaseData, node_name: str):
+    def _run(self, g: Graph, _input: BaseData, node_name: str):
         print(f"---- Starting node {node_name}")
         print(f'node_name {node_name}')
 

{indexify-0.0.39 → indexify-0.0.40}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "indexify"
-version = "0.0.39"
+version = "0.0.40"
 description = "Python Client for Indexify"
 authors = ["Diptanu Gon Choudhury <diptanuc@gmail.com>", "Lucas Jackson <lucas@tensorlake.ai>", "Vijay Parthasarathy <vijay2win@gmail.com>"]
 license = "Apache 2.0"
@@ -13,6 +13,7 @@ python = "^3.9"
 httpx = { version = "^0", extras = ["http2"] }
 pyyaml = "^6"
 pydantic = "^2.8"
+cloudpickle = "^3"
 
 [tool.poetry.dev-dependencies]
 black = "^22.3.0"

indexify-0.0.39/indexify/graph.py

@@ -1,23 +0,0 @@
-from .extractor_sdk import extractor, Extractor
-
-from typing import Type, Union
-from pydantic import BaseModel
-
-from .run_graph import RunGraph
-from .local_runner import LocalRunner
-
-
-def Graph(
-    name: str,
-    input: Type[BaseModel],
-    start_node: Union[extractor, Extractor],
-    run_local: bool,
-) -> RunGraph:
-
-    if run_local:
-        runner = LocalRunner()
-    else:
-        raise NotImplementedError("Remote runner not supported yet")
-
-    graph = RunGraph(name=name, input=input, start_node=start_node, runner=runner)
-    return graph