text: __copyright__ = "Copyright (c) 2020 Jina AI Limited. All rights reserved."
  __license__ = "Apache-2.0"
  
  import os
  import itertools as it
  
  import click
  from jina.flow import Flow
  from jina import Document
  
  
  def config():
      os.environ["JINA_DATA_FILE"] = os.environ.get(
          "JINA_DATA_FILE", "data/character-lines.csv"
      )
      os.environ["JINA_WORKSPACE"] = os.environ.get("JINA_WORKSPACE", "workspace")
  
      os.environ["JINA_PORT"] = os.environ.get("JINA_PORT", str(45678))
  
  
  def print_topk(resp, sentence):
      for d in resp.search.docs:
          print(f"Ta-Dah🔮, here are what we found for: {sentence}")
          for idx, match in enumerate(d.matches):
  
              score = match.score.value
              if score < 0.0:
                  continue
              character = match.tags['character']
              dialog = match.text.strip()
              print(f'> {idx:>2d}({score:.2f}). {character.upper()} said, "{dialog}"')
  
  
  def index_generator(filepath: str, num_docs: int):
      def sample(iterable):
          for i in iterable:
              yield i
  
      with open(filepath, 'r') as f:
          for line in it.islice(sample(f), num_docs):
              character, sentence = line.split('[SEP]')
              document = Document()
              document.text = sentence
              document.tags['character'] = character
              yield document
  
  
  def index(num_docs):
      f = Flow().load_config("flow-index.yml")
  
      with f:
          f.index(input_fn=index_generator(filepath=os.environ["JINA_DATA_FILE"], num_docs=num_docs),
                  batch_size=8)
  
  
  def query(top_k):
      f = Flow().load_config("flow-query.yml")
      with f:
          while True:
              text = input("please type a sentence: ")
              if not text:
                  break
  
              def ppr(x):
                  print_topk(x, text)
  
              f.search_lines(lines=[text, ], output_fn=ppr, top_k=top_k)
  
  
  def query_restful():
      f = Flow().load_config("flow-query.yml")
      f.use_rest_gateway()
      with f:
          f.block()
  
  
  def dryrun():
      f = Flow().load_config("flow-index.yml")
      with f:
          f.dry_run()
  
  
  @click.command()
  @click.option(
      "--task",
      "-t",
      type=click.Choice(
          ["index", "query", "query_restful", "dryrun"], case_sensitive=False
      ),
  )
  @click.option("--num_docs", "-n", default=50)
  @click.option("--top_k", "-k", default=5)
  def main(task, num_docs, top_k):
      config()
      workspace = os.environ["JINA_WORKSPACE"]
      if task == "index":
          if os.path.exists(workspace):
              print(f'\n +----------------------------------------------------------------------------------+ \
                      \n |                                   🤖🤖🤖                                         | \
                      \n | The directory {workspace} already exists. Please remove it before indexing again.  | \
                      \n |                                   🤖🤖🤖                                         | \
                      \n +----------------------------------------------------------------------------------+')
          index(num_docs)
      if task == "query":
          if not os.path.exists(workspace):
              print(f"The directory {workspace} does not exist. Please index first via `python app.py -t index`")
          query(top_k)
      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`")
          query_restful()
      if task == "dryrun":
          dryrun()
  
  
  if __name__ == "__main__":
      main()
  