Browse Source
Update code for dc2
Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
2 changed files with
12 additions and
27 deletions
-
README.md
-
faiss.py
|
|
@ -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') |
|
|
|
``` |
|
|
|
|
|
|
|
<img src="./result.png" height="100px"/> |
|
|
|
|
|
@ -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 |
|
|
|