diff --git a/README.md b/README.md index d704002..7cb4061 100644 --- a/README.md +++ b/README.md @@ -16,37 +16,21 @@ Search embedding in [Faiss](https://github.com/facebookresearch/faiss), **please ## Code Example -### Insert data into Faiss first - -```python -import numpy as np -import towhee - -vec = np.random.random((10, 100)).astype('float32') -ids = list(i for i in range(10)) - -x = towhee.dc['id'](ids) \ - .runas_op['id', 'vec'](func=lambda x: vec[x]) \ - .to_faiss['id', 'vec'](findex='index.bin') -``` +> please insert data into faiss first. ### Example -*Write the pipeline in simplified style:* +*Write a same pipeline with explicit inputs/outputs name specifications:* ```python -query = vec[0:2] -towhee.dc(query) \ - .ann_search.faiss(findex='index.bin') -``` +from towhee.dc2 import pipe, ops -*Write a same pipeline with explicit inputs/outputs name specifications:* +p = pipe.input('vec') \ + .flat_map('vec', 'rows', ops.ann_search.faiss(findex='index.bin')) \ + .map('rows', ('id', 'score'), lambda x: (x[0], x[1], x[2])) \ + .output('id', 'score') -```python -query = vec[0:2] -towhee.dc['vec'](query) \ - .ann_search.faiss['vec', 'results'](findex='index.bin') \ - .show() +p('cat') ``` diff --git a/faiss.py b/faiss.py index fc9569c..b222a25 100644 --- a/faiss.py +++ b/faiss.py @@ -2,8 +2,8 @@ import numpy as np from pathlib import Path import faiss from towhee import register -from towhee.utils.faiss_utils import KVStorage -from towhee.functional.entity import Entity +from towhee.utils.thirdparty.faiss_utils import KVStorage +# from towhee.functional.entity import Entity @register(output_schema=['result']) @@ -56,5 +56,6 @@ class Faiss: k = self.kv_storage.get(ids[i]) else: k = ids[i] - result.append(Entity(**{'key': k, 'score': scores[0][i]})) + result.extend([k, scores[0][i]]) + # result.append(Entity(**{'key': k, 'score': scores[0][i]})) return result