diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..db8a3b1 --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ +from .milvus import Milvus + +def milvus(*args, **kwargs): + return Milvus(*args, **kwargs) \ No newline at end of file diff --git a/milvus.py b/milvus.py new file mode 100644 index 0000000..7c60bd2 --- /dev/null +++ b/milvus.py @@ -0,0 +1,38 @@ +import numpy as np +from towhee import register +from pymilvus import Collection + + +@register(output_schema=['mr']) +class Milvus: + """ + Milvus ANN index class. + """ + + def __init__(self, collection): + """ + Get an existing collection. + """ + if isinstance(collection, str): + collection = Collection(collection) + self._collection = collection + + def __call__(self, data): + """ + Insert data to Milvus. + + Args: + data (`list`): + The data to insert into milvus. + + Returns: + A MutationResult object contains `insert_count` represents how many and a `primary_keys` of primary keys. + + """ + vectors = [] + if isinstance(data, np.ndarray): + data = [data] + for v in data: + vectors.append(v if isinstance(v, list) else [v]) + mr = self._collection.insert(vectors) + return mr diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..9c1554a --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +pymilvus +towhee +numpy \ No newline at end of file diff --git a/result.png b/result.png new file mode 100644 index 0000000..f39e183 Binary files /dev/null and b/result.png differ