Frontend start command.

backend/app.py

@@ -1,20 +1,32 @@
+__copyright__ = "Copyright (c) 2021 Jina AI Limited. All rights reserved."
+__license__ = "Apache-2.0"
+
 import os
 import itertools
-from pprint import pprint
-from jina import Flow, Document, DocumentArray
-from jina.types.arrays.memmap import DocumentArrayMemmap
-from jina.parsers.helloworld import set_hw_chatbot_parser
 import csv
-from backend_config import backend_port, backend_workdir, backend_datafile, text_length, max_docs
-from executors import MyTransformer, MyIndexer
+import shutil
+import click
+import sys
+from backend_config import (
+    text_length,
+    max_docs,
+    backend_datafile,
+    backend_port,
+    backend_workdir,
+)
+from executors import MyTransformer, DiskIndexer
+
+from jina import Flow, Document
 
 try:
     __import__("pretty_errors")
 except ImportError:
     pass
 
 
-def trim_string(input_string: str, word_count: int = text_length, sep: str = " ") -> str:
+def trim_string(
+    input_string: str, word_count: int = text_length, sep: str = " "
+) -> str:
     """
     Trim a string to a certain number of words.
     :param input_string: string to trim
@@ -46,48 +58,67 @@ def prep_docs(input_file: str, max_docs=max_docs):
             yield doc
 
 
-def run_appstore_flow(inputs, args) -> None:
-    """
-    Execute the app store example. Indexes data and presents REST endpoint
-    :param inputs: Documents or DocumentArrays to input
-    :args: arguments like port, workdir, etc
-    :return: None
-    """
-
-    # Create Flow and add
-    #   - MyTransformer (an encoder Executor)
-    #   - MyIndexer (a simple indexer Executor)
+def index():
     flow = (
         Flow()
-        .add(uses=MyTransformer, parallel=args.parallel)
-        # .add(uses=EmbeddingIndexer)
-        # .add(uses=KeyValueIndexer)
-        .add(uses=MyIndexer, workspace=args.workdir)
+        .add(uses=MyTransformer, parallel=2, name="encoder")
+        .add(uses=DiskIndexer, workspace=backend_workdir, name="indexer")
     )
-    # flow = Flow.load_config('flows/index.yml')
 
-    # Open the Flow
     with flow:
-        # Start index pipeline, taking inputs then printing the processed DocumentArray
-        flow.post(on="/index", inputs=inputs)
+        flow.post(
+            on="/index",
+            inputs=prep_docs(input_file=backend_datafile, max_docs=max_docs),
+            request_size=64,
+            read_mode="r",
+        )
 
-        # Start REST gateway so clients can query via Streamlit or other frontend (like Jina Box)
-        flow.use_rest_gateway(backend_port)
-
-        # Block the process to keep it open. Otherwise it will just close and no-one could connect
-        flow.block()
 
+def query_restful():
+    flow = (
+        Flow()
+        .add(uses=MyTransformer, name="encoder")
+        .add(uses=DiskIndexer, workspace=backend_workdir, name="indexer")
+    )
 
-if __name__ == "__main__":
+    with flow:
+        flow.protocol = "http"
+        flow.port_expose = backend_port
+        flow.block()
 
-    # Get chatbot's default arguments
-    args = set_hw_chatbot_parser().parse_args()
 
-    # Change a few things
-    args.workdir = backend_workdir
+@click.command()
+@click.option(
+    "--task",
+    "-t",
+    type=click.Choice(["index", "query_restful"], case_sensitive=False),
+)
+@click.option("--num_docs", "-n", default=max_docs)
+@click.option("--force", "-f", is_flag=True)
+def main(task: str, num_docs: int, force: bool):
+    workspace = backend_workdir
+    if task == "index":
+        if os.path.exists(workspace):
+            if force:
+                shutil.rmtree(workspace)
+            else:
+                print(
+                    f"\n +----------------------------------------------------------------------------------+ \
+                        \n |                                   🤖🤖🤖                                         | \
+                        \n | The directory {workspace} already exists. Please remove it before indexing again.  | \
+                        \n |                                   🤖🤖🤖                                         | \
+                        \n +----------------------------------------------------------------------------------+"
+                )
+                sys.exit(1)
+        index()
+    if task == "query_restful":
+        if not os.path.exists(workspace):
+            print(
+                f"The directory {workspace} does not exist. Please index first via `python app.py -t index`"
+            )
+            sys.exit(1)
+        query_restful()
 
-    # Convert the csv file to a DocumentArray
-    docs = prep_docs(input_file=backend_datafile)
 
-    # Run the Flow
-    run_appstore_flow(inputs=docs, args=args)
+if __name__ == "__main__":
+    main()

backend/backend_config.py

@@ -7,7 +7,7 @@
 backend_workdir = "workspace"
 backend_datafile = "./data/appstore_games-shuffled.csv"
 text_length = 50 # How many words to index for each app? Longer = more accurate, shorter = quicker
-max_docs = 5000 # How many apps to index
+max_docs = 3000 # How many apps to index
 
 # dataset
 dataset_url = "https://github.com/alexcg1/ml-datasets/blob/master/nlp/strategy_games/appstore_games.csv?raw=true" 

backend/executors.py

@@ -2,6 +2,7 @@
 from backend_config import backend_model, backend_top_k
 
 import numpy as np
+import os
 import torch
 from transformers import AutoModel, AutoTokenizer
 
@@ -89,17 +90,67 @@ def encode(self, docs: "DocumentArray", *args, **kwargs):
                 doc.embedding = embed
 
 
-class MyIndexer(Executor):
+def _get_ones(x, y):
+    return np.ones((x, y))
+
+
+def _ext_A(A):
+    nA, dim = A.shape
+    A_ext = _get_ones(nA, dim * 3)
+    A_ext[:, dim : 2 * dim] = A
+    A_ext[:, 2 * dim :] = A ** 2
+    return A_ext
+
+
+def _ext_B(B):
+    nB, dim = B.shape
+    B_ext = _get_ones(dim * 3, nB)
+    B_ext[:dim] = (B ** 2).T
+    B_ext[dim : 2 * dim] = -2.0 * B.T
+    del B
+    return B_ext
+
+
+def _euclidean(A_ext, B_ext):
+    sqdist = A_ext.dot(B_ext).clip(min=0)
+    return np.sqrt(sqdist)
+
+
+def _norm(A):
+    return A / np.linalg.norm(A, ord=2, axis=1, keepdims=True)
+
+
+def _cosine(A_norm_ext, B_norm_ext):
+    return A_norm_ext.dot(B_norm_ext).clip(min=0) / 2
+
+
+
+class DiskIndexer(Executor):
     """Simple indexer class """
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self._docs = DocumentArray()
         self.top_k = backend_top_k
+        if os.path.exists(self.save_path):
+            self._docs = DocumentArray.load(self.save_path)
+        else:
+            self._docs = DocumentArray()
+
+    @property
+    def save_path(self):
+        if not os.path.exists(self.workspace):
+            os.makedirs(self.workspace)
+        return os.path.join(self.workspace, 'apps.json')
+
+    def close(self):
+        self._docs.save(self.save_path)
+
 
     @requests(on="/index")
     def index(self, docs: "DocumentArray", **kwargs):
         self._docs.extend(docs)
+        return docs
 
     @requests(on="/search")
     def search(self, docs: "DocumentArray", **kwargs):
@@ -112,8 +163,9 @@ def search(self, docs: "DocumentArray", **kwargs):
         for _q, _ids, _dists in zip(docs, idx, dist):
             for _id, _dist in zip(_ids, _dists):
                 d = Document(self._docs[int(_id)], copy=True)
-                d.score.value = 1 - _dist
+                # d.score.value = 1 - _dist
                 _q.matches.append(d)
+        return docs
 
     @staticmethod
     def _get_sorted_top_k(
@@ -130,37 +182,3 @@ def _get_sorted_top_k(
             dist = np.take_along_axis(dist, idx_fs, axis=1)
 
         return idx, dist
-
-
-def _get_ones(x, y):
-    return np.ones((x, y))
-
-
-def _ext_A(A):
-    nA, dim = A.shape
-    A_ext = _get_ones(nA, dim * 3)
-    A_ext[:, dim : 2 * dim] = A
-    A_ext[:, 2 * dim :] = A ** 2
-    return A_ext
-
-
-def _ext_B(B):
-    nB, dim = B.shape
-    B_ext = _get_ones(dim * 3, nB)
-    B_ext[:dim] = (B ** 2).T
-    B_ext[dim : 2 * dim] = -2.0 * B.T
-    del B
-    return B_ext
-
-
-def _euclidean(A_ext, B_ext):
-    sqdist = A_ext.dot(B_ext).clip(min=0)
-    return np.sqrt(sqdist)
-
-
-def _norm(A):
-    return A / np.linalg.norm(A, ord=2, axis=1, keepdims=True)
-
-
-def _cosine(A_norm_ext, B_norm_ext):
-    return A_norm_ext.dot(B_norm_ext).clip(min=0) / 2