code: """
  View stage example selector.
  
  | Copyright 2017-2023, Voxel51, Inc.
  | `voxel51.com <https://voxel51.com/>`_
  |
  """
  import os
  import pickle
  
  from langchain.prompts import FewShotPromptTemplate, PromptTemplate
  import numpy as np
  import pandas as pd
  from scipy.spatial.distance import cosine
  
  # pylint: disable=relative-beyond-top-level
  from .utils import get_embedding_function, get_cache, hash_query
  
  
  ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  EXAMPLES_DIR = os.path.join(ROOT_DIR, "examples")
  
  EXAMPLE_EMBEDDINGS_PATH = os.path.join(
      EXAMPLES_DIR, "viewstage_embeddings.pkl"
  )
  VIEW_STAGE_EXAMPLES_PATH = os.path.join(EXAMPLES_DIR, "viewstage_examples.csv")
  
  VIEW_STAGE_EXAMPLE_PROMPT = PromptTemplate(
      input_variables=["input", "output"],
      template="Input: {input}\nOutput: {output}",
  )
  
  
  def get_or_create_embeddings(queries):
      if os.path.isfile(EXAMPLE_EMBEDDINGS_PATH):
          with open(EXAMPLE_EMBEDDINGS_PATH, "rb") as f:
              example_embeddings = pickle.load(f)
      else:
          example_embeddings = {}
  
      query_hashes = []
      new_hashes = []
      new_queries = []
  
      for query in queries:
          key = hash_query(query)
          query_hashes.append(key)
  
          if key not in example_embeddings:
              new_hashes.append(key)
              new_queries.append(query)
  
      if new_queries:
          print("Generating %d embeddings..." % len(new_queries))
          model = get_embedding_function()
          new_embeddings = model(new_queries)
          for key, embedding in zip(new_hashes, new_embeddings):
              example_embeddings[key] = embedding
  
      if new_queries:
          print("Saving embeddings to disk...")
  
          with open(EXAMPLE_EMBEDDINGS_PATH, "wb") as f:
              pickle.dump(example_embeddings, f)
  
      return example_embeddings
  
  
  def has_geo_field(sample_collection):
      types = list(sample_collection.get_field_schema(flat=True).values())
      types = [type(t) for t in types]
      return any(["Geo" in t.__name__ for t in types])
  
  
  def get_label_type(sample_collection, field_name):
      sample = sample_collection.first()
      field = sample.get_field(field_name)
      field_type = str(type(field).__name__).lower()
      field_type = field_type[:-1] if field_type.endswith("s") else field_type
      return field_type
  
  
  def _replace_run_keys(prompt, runs):
      if "text_similarity" in runs:
          prompt = prompt.replace("TEXT_SIM_KEY", runs["text_similarity"]["key"])
      if "image_similarity" in runs:
          prompt = prompt.replace(
              "IMAGE_SIM_KEY", runs["image_similarity"]["key"]
          )
      if "evaluation" in runs:
          prompt = prompt.replace("EVAL_KEY", runs["evaluation"]["key"])
      if "uniqueness" in runs:
          prompt = prompt.replace(
              "UNIQUENESS_FIELD", runs["uniqueness"]["uniqueness_field"]
          )
      return prompt
  
  
  def _count_empty_class_names(label_field):
      return [list(class_name.values())[0] for class_name in label_field].count(
          []
      )
  
  
  def _reduce_label_fields(label_fields):
      label_field_keys = list(label_fields.keys())
      if len(label_field_keys) == 0:
          return None, None
      elif len(label_field_keys) > 0:
          empty_counts = [
              _count_empty_class_names(label_fields[key])
              for key in label_field_keys
          ]
          min_empty_count = min(empty_counts)
          valid_keys = [
              key
              for key, count in zip(label_field_keys, empty_counts)
              if count == min_empty_count
          ]
          return {key: label_fields[key] for key in valid_keys}, min_empty_count
  
  
  def _parse_runs_and_labels(runs, label_fields):
      reduced_label_fields, count = _reduce_label_fields(label_fields.copy())
      reduced_runs = runs.copy()
      if count is not None and count > 0 and "text_similarity" in reduced_runs:
          reduced_label_fields = None
  
      return reduced_runs, reduced_label_fields
  
  
  def _get_evaluation_type(sample_collection, eval_key):
      eval_cls = sample_collection.get_evaluation_info(eval_key).config.cls
      if "openimages" in eval_cls:
          return "detection"
      elif "coco" in eval_cls:
          return "detection"
      elif "activitynet" in eval_cls:
          return "detection"
      elif "classification" in eval_cls:
          return "classification"
      return None
  
  
  def _load_examples():
      examples = pd.read_csv(VIEW_STAGE_EXAMPLES_PATH, on_bad_lines="skip")
      examples["meta"] = examples["metadata"]
      examples["contains_match"] = examples["stages"].str.contains("match\(")
      examples["contains_filter_labels"] = examples["stages"].str.contains(
          "filter_labels\("
      )
      examples["mfl"] = (
          examples["contains_match"] | examples["contains_filter_labels"]
      )
      examples["hash"] = examples["query"].apply(lambda x: hash_query(x))
  
      queries = examples["query"].tolist()
      embeddings = get_or_create_embeddings(queries)
  
      embeddings = {
          key: np.array(embeddings[key]) for key in examples["hash"].tolist()
      }
      return examples, embeddings
  
  
  def get_examples():
      cache = get_cache()
      keys = ("viewstage_examples", "viewstage_embeddings")
      if keys[0] not in cache or keys[1] not in cache:
          cache[keys[0]], cache[keys[1]] = _load_examples()
  
      return cache[keys[0]], cache[keys[1]]
  
  
  def _get_filtered_examples(sample_collection, runs, label_fields):
      examples, embeddings = get_examples()
      media_type = sample_collection.media_type
      _filter = examples["media_type"].isin([media_type, "all"])
  
      red_runs, red_label_fields = _parse_runs_and_labels(runs, label_fields)
      geo = has_geo_field(sample_collection)
      text_sim = "text_similarity" in red_runs
      image_sim = "image_similarity" in red_runs
      meta = "metadata" in red_runs
      eval = "evaluation" in red_runs
  
      if red_label_fields or eval:
          if red_label_fields:
              label_field_types = list(
                  set(
                      [
                          get_label_type(sample_collection, field)
                          for field in red_label_fields
                      ]
                  )
              )
          else:
              label_field_types = []
  
          if eval:
              eval_key = red_runs["evaluation"]["key"]
              eval_types = [_get_evaluation_type(sample_collection, eval_key)]
          else:
              eval_types = []
  
          label_types = list(set(label_field_types + eval_types + ["all"]))
          _filter = _filter & examples["label_type"].isin(label_types)
  
      ## contains match() or filter_labels() in stages
      mfl_cond = red_label_fields and not text_sim
  
      conds = [geo, text_sim, image_sim, meta, eval, mfl_cond]
      strs = ["geo", "text_sim", "image_sim", "meta", "eval", "mfl"]
  
      for cond, cond_str in zip(conds, strs):
          if not cond:
              _filter = _filter & (examples[cond_str] == False)
  
      filtered_examples = examples[_filter]
  
      filtered_queries, filtered_stages, hashes = (
          filtered_examples["query"].tolist(),
          filtered_examples["stages"].tolist(),
          filtered_examples["hash"].tolist(),
      )
      filtered_embeddings = [embeddings[key] for key in hashes]
  
      return filtered_queries, filtered_stages, filtered_embeddings
  
  
  def get_similar_examples(sample_collection, query, runs, label_fields):
      ex_queries, ex_stages, ex_embeddings = _get_filtered_examples(
          sample_collection, runs, label_fields
      )
  
      model = get_embedding_function()
      query_embedding = np.array(model([query]))
      if len(query_embedding.shape) == 2:
          query_embedding = query_embedding[0]
  
      dists = np.array([cosine(query_embedding, emb) for emb in ex_embeddings])
  
      sorted_ix = np.argsort(dists).astype(int)
  
      k = 20
      similar_queries = [ex_queries[ix] for ix in sorted_ix[:k]]
      similar_stages = [ex_stages[ix] for ix in sorted_ix[:k]]
  
      return [
          {"input": sq, "output": ss}
          for sq, ss in zip(similar_queries, similar_stages)
      ]
  
  
  def generate_view_stage_examples_prompt_template(
      sample_collection, query, runs, label_fields
  ):
      examples = get_similar_examples(
          sample_collection, query, runs, label_fields
      )
      example_prompt = VIEW_STAGE_EXAMPLE_PROMPT
      return FewShotPromptTemplate(
          examples=examples,
          example_prompt=example_prompt,
          prefix="Generate code to produce the FiftyOne view stages for the following prompts:\n",
          suffix="Input: {text}\nOutput:",
          input_variables=["text"],
      )
  
  
  def generate_view_stage_examples_prompt(
      sample_collection, query, runs, label_fields
  ):
      similar_examples_prompt_template = (
          generate_view_stage_examples_prompt_template(
              sample_collection, query, runs, label_fields
          )
      )
  
      prompt = similar_examples_prompt_template.format(text=query)
      return _replace_run_keys(prompt, runs)
  
apis: langchain.prompts.PromptTemplate langchain.prompts.FewShotPromptTemplate
extract_api: [((489, 523), 'os.path.join', 'os.path.join', (['ROOT_DIR', '"""examples"""'], {}), "(ROOT_DIR, 'examples')\n", (501, 523), False, 'import os\n'), ((551, 605), 'os.path.join', 'os.path.join', (['EXAMPLES_DIR', '"""viewstage_embeddings.pkl"""'], {}), "(EXAMPLES_DIR, 'viewstage_embeddings.pkl')\n", (563, 605), False, 'import os\n'), ((639, 691), 'os.path.join', 'os.path.join', (['EXAMPLES_DIR', '"""viewstage_examples.csv"""'], {}), "(EXAMPLES_DIR, 'viewstage_examples.csv')\n", (651, 691), False, 'import os\n'), ((721, 825), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['input', 'output']", 'template': '"""Input: {input}\nOutput: {output}"""'}), '(input_variables=[\'input\', \'output\'], template=\n    """Input: {input}\nOutput: {output}""")\n', (735, 825), False, 'from langchain.prompts import FewShotPromptTemplate, PromptTemplate\n'), ((877, 916), 'os.path.isfile', 'os.path.isfile', (['EXAMPLE_EMBEDDINGS_PATH'], {}), '(EXAMPLE_EMBEDDINGS_PATH)\n', (891, 916), False, 'import os\n'), ((4348, 4406), 'pandas.read_csv', 'pd.read_csv', (['VIEW_STAGE_EXAMPLES_PATH'], {'on_bad_lines': '"""skip"""'}), "(VIEW_STAGE_EXAMPLES_PATH, on_bad_lines='skip')\n", (4359, 4406), True, 'import pandas as pd\n'), ((8216, 8455), 'langchain.prompts.FewShotPromptTemplate', 'FewShotPromptTemplate', ([], {'examples': 'examples', 'example_prompt': 'example_prompt', 'prefix': '"""Generate code to produce the FiftyOne view stages for the following prompts:\n"""', 'suffix': '"""Input: {text}\nOutput:"""', 'input_variables': "['text']"}), '(examples=examples, example_prompt=example_prompt,\n    prefix=\n    """Generate code to produce the FiftyOne view stages for the following prompts:\n"""\n    , suffix="""Input: {text}\nOutput:""", input_variables=[\'text\'])\n', (8237, 8455), False, 'from langchain.prompts import FewShotPromptTemplate, PromptTemplate\n'), ((446, 471), 'os.path.abspath', 'os.path.abspath', (['__file__'], {}), '(__file__)\n', (461, 471), False, 'import os\n'), ((4933, 4958), 'numpy.array', 'np.array', (['embeddings[key]'], {}), '(embeddings[key])\n', (4941, 4958), True, 'import numpy as np\n'), ((1006, 1020), 'pickle.load', 'pickle.load', (['f'], {}), '(f)\n', (1017, 1020), False, 'import pickle\n'), ((1751, 1785), 'pickle.dump', 'pickle.dump', (['example_embeddings', 'f'], {}), '(example_embeddings, f)\n', (1762, 1785), False, 'import pickle\n'), ((7604, 7632), 'scipy.spatial.distance.cosine', 'cosine', (['query_embedding', 'emb'], {}), '(query_embedding, emb)\n', (7610, 7632), False, 'from scipy.spatial.distance import cosine\n'), ((7677, 7694), 'numpy.argsort', 'np.argsort', (['dists'], {}), '(dists)\n', (7687, 7694), True, 'import numpy as np\n')]