Add Elasticsearch example (#68)

This closes #67.

Author: RussellLuo

examples/es_example/README.md

@@ -0,0 +1,47 @@
+# Elasticsearch
+
+This example illustrates how to use RocketQA along with [Elasticsearch](https://www.elastic.co/).
+
+
+## Prerequisites
+
+### Install Dependencies
+
+```console
+$ python3 -m venv venv
+$ source venv/bin/activate
+$ pip3 install -r requirements.txt
+```
+
+### Run Elasticsearch
+
+Run Elasticsearch in development mode:
+
+```console
+$ docker run -d --name elasticsearch -p 9200:9200 -p 9300:9300 -e "discovery.type=single-node" -e "ELASTIC_PASSWORD=123456" elasticsearch:8.4.2
+```
+
+## Usage
+
+### Index
+
+Prepare the data (stored at `data/test.tsv`) in the following format:
+
+```
+title_1\tparagraph_1\n
+title_2\tparagraph_2\n
+...
+```
+
+Create the index and save the data into the index:
+
+```console
+$ curl -XPUT -u elastic:123456 -k -H "Content-Type: application/json" https://localhost:9200/test-index -d @mappings.json
+$ python3 index.py zh data/test.tsv test-index
+```
+
+### Query
+
+```console
+$ python3 query.py
+```

examples/es_example/index.py

@@ -0,0 +1,72 @@
+# -*- coding: utf-8 -*-
+
+import argparse
+import os
+import sys
+
+import faiss
+import numpy as np
+import rocketqa
+from elasticsearch import Elasticsearch, helpers
+
+
+class Indexer:
+    def __init__(self, es_client, index_name, model):
+        self.es_client = es_client
+        self.index_name = index_name
+        self.dual_encoder = rocketqa.load_model(
+            model=model,
+            use_cuda=False, # GPU: True
+            device_id=0,
+            batch_size=32,
+        )
+
+    def index(self, tps):
+        titles, paras = zip(*tps)
+        embs = self.dual_encoder.encode_para(para=paras, title=titles)
+    
+        def gen_actions():
+            for i, emb in enumerate(embs):
+                # Normalize the NumPy array to a unit vector to use `dot_product` similarity,
+                # see https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html#dense-vector-params.
+                emb = emb / np.linalg.norm(emb)
+                yield dict(
+                    _index=self.index_name,
+                    _id=i+1,
+                    _source=dict(
+                        title=titles[i],
+                        paragraph=paras[i],
+                        vector=emb,
+                    ),
+                )
+        return helpers.bulk(self.es_client, gen_actions())
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('lang', choices=['zh', 'en'], help='The language')
+    parser.add_argument('data_file', help='The data file')
+    parser.add_argument('index_name', help='The index name')
+    args = parser.parse_args()
+
+    if args.lang == 'zh':
+        model = 'zh_dureader_de_v2'
+    elif args.lang == 'en':
+        model = 'v1_marco_de'
+
+    with open(args.data_file) as f:
+      tps = [line.strip().split('\t') for line in f]
+
+    es_client = Elasticsearch(
+        "https://localhost:9200",
+        http_auth=("elastic", "123456"),
+        verify_certs=False,
+    )
+
+    indexer = Indexer(es_client, args.index_name, model)
+    result = indexer.index(tps)
+    print(result)
+
+
+if __name__ == '__main__':
+    main()

examples/es_example/mappings.json

@@ -0,0 +1,23 @@
+{
+  "mappings": {
+    "_source": {
+      "excludes": [
+        "vector"
+      ]
+    },
+    "properties": {
+      "vector": {
+        "type": "dense_vector",
+        "dims": 768,
+        "index": true,
+        "similarity": "dot_product"
+      },
+      "title": {
+        "type": "text"
+      },
+      "paragraph": {
+        "type": "text"
+      }
+    }
+  }
+}

examples/es_example/query.py

@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+
+import sys
+import time
+
+import numpy as np
+import rocketqa
+from elasticsearch import Elasticsearch
+
+
+class Querier:
+    def __init__(self, es_client, index_name, de_model, ce_model):
+        self.es_client = es_client
+        self.index_name = index_name
+        self.dual_encoder = rocketqa.load_model(
+            model=de_model,
+            use_cuda=False, # GPU: True
+            device_id=0,
+            batch_size=32,
+        )
+        self.cross_encoder = rocketqa.load_model(
+            model=ce_model,
+            use_cuda=False, # GPU: True
+            device_id=0,
+            batch_size=32,
+        )
+
+    def encode(self, query):
+        embs = self.dual_encoder.encode_query(query=[query])
+        vector = list(embs)[0]
+        # Normalize the NumPy array to a unit vector to use `dot_product` similarity,
+        # see https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html#dense-vector-params.
+        vector = vector / np.linalg.norm(vector)
+        return vector
+
+    def search(self, query, topk=10):
+        vector = self.encode(query)
+        knn = dict(
+            field="vector",
+            query_vector=vector,
+            k=topk,
+            num_candidates=100,
+        )
+        result = self.es_client.knn_search(index=self.index_name, knn=knn)
+
+        candidates = [
+            dict(
+                title=doc['_source']['title'],
+                para=doc['_source']['paragraph'],
+            )
+            for doc in result['hits']['hits']
+        ]
+        return candidates
+
+    def sort(self, query, candidates):
+        queries = [query] * len(candidates)
+        titles = [c['title'] for c in candidates]
+        paras = [c['para'] for c in candidates]
+        ranking_score = self.cross_encoder.matching(query=queries, para=paras, title=titles)
+    
+        answers = [
+            dict(
+                title=titles[i],
+                para=paras[i],
+                score=score,
+            )
+            for i, score in enumerate(ranking_score)
+        ]
+        return sorted(answers, key=lambda a: a['score'], reverse=True)
+
+
+def main():
+    es_client = Elasticsearch(
+        "https://localhost:9200",
+        http_auth=("elastic", "123456"),
+        verify_certs=False,
+    )
+    querier = Querier(es_client, "test-index", 'zh_dureader_de_v2', 'zh_dureader_ce_v2')
+
+    while True:
+        query = input('Query: ')
+
+        candidates = querier.search(query)
+        print('Candidates:')
+        for c in candidates:
+            print(c['title'], '\t', c['para'])
+
+        answers = querier.sort(query, candidates)
+        print('Answers:')
+        for a in answers:
+            print(a['title'], '\t', a['para'], '\t', a['score'])
+
+
+if __name__ == '__main__':
+    main()

examples/es_example/requirements.txt

@@ -0,0 +1,4 @@
+paddlepaddle==2.3.2
+rocketqa==1.1.0
+elasticsearch==8.5.0
+numpy==1.21.6