code: """General utils functions."""
  
  import asyncio
  import os
  import random
  import sys
  import time
  import traceback
  import uuid
  from contextlib import contextmanager
  from dataclasses import dataclass
  from functools import partial, wraps
  from itertools import islice
  from pathlib import Path
  from typing import (
      Any,
      AsyncGenerator,
      Callable,
      Dict,
      Generator,
      Iterable,
      List,
      Optional,
      Protocol,
      Set,
      Type,
      Union,
      cast,
      runtime_checkable,
  )
  
  
  class GlobalsHelper:
      """Helper to retrieve globals.
  
      Helpful for global caching of certain variables that can be expensive to load.
      (e.g. tokenization)
  
      """
  
      _tokenizer: Optional[Callable[[str], List]] = None
      _stopwords: Optional[List[str]] = None
  
      @property
      def tokenizer(self) -> Callable[[str], List]:
          """Get tokenizer. TODO: Deprecated."""
          if self._tokenizer is None:
              tiktoken_import_err = (
                  "`tiktoken` package not found, please run `pip install tiktoken`"
              )
              try:
                  import tiktoken
              except ImportError:
                  raise ImportError(tiktoken_import_err)
              enc = tiktoken.get_encoding("gpt2")
              self._tokenizer = cast(Callable[[str], List], enc.encode)
              self._tokenizer = partial(self._tokenizer, allowed_special="all")
          return self._tokenizer  # type: ignore
  
      @property
      def stopwords(self) -> List[str]:
          """Get stopwords."""
          if self._stopwords is None:
              try:
                  import nltk
                  from nltk.corpus import stopwords
              except ImportError:
                  raise ImportError(
                      "`nltk` package not found, please run `pip install nltk`"
                  )
  
              from llama_index.utils import get_cache_dir
  
              cache_dir = get_cache_dir()
              nltk_data_dir = os.environ.get("NLTK_DATA", cache_dir)
  
              # update nltk path for nltk so that it finds the data
              if nltk_data_dir not in nltk.data.path:
                  nltk.data.path.append(nltk_data_dir)
  
              try:
                  nltk.data.find("corpora/stopwords")
              except LookupError:
                  nltk.download("stopwords", download_dir=nltk_data_dir)
              self._stopwords = stopwords.words("english")
          return self._stopwords
  
  
  globals_helper = GlobalsHelper()
  
  
  # Global Tokenizer
  @runtime_checkable
  class Tokenizer(Protocol):
      def encode(self, text: str, *args: Any, **kwargs: Any) -> List[Any]:
          ...
  
  
  def set_global_tokenizer(tokenizer: Union[Tokenizer, Callable[[str], list]]) -> None:
      import llama_index
  
      if isinstance(tokenizer, Tokenizer):
          llama_index.global_tokenizer = tokenizer.encode
      else:
          llama_index.global_tokenizer = tokenizer
  
  
  def get_tokenizer() -> Callable[[str], List]:
      import llama_index
  
      if llama_index.global_tokenizer is None:
          tiktoken_import_err = (
              "`tiktoken` package not found, please run `pip install tiktoken`"
          )
          try:
              import tiktoken
          except ImportError:
              raise ImportError(tiktoken_import_err)
          enc = tiktoken.encoding_for_model("gpt-3.5-turbo")
          tokenizer = partial(enc.encode, allowed_special="all")
          set_global_tokenizer(tokenizer)
  
      assert llama_index.global_tokenizer is not None
      return llama_index.global_tokenizer
  
  
  def get_new_id(d: Set) -> str:
      """Get a new ID."""
      while True:
          new_id = str(uuid.uuid4())
          if new_id not in d:
              break
      return new_id
  
  
  def get_new_int_id(d: Set) -> int:
      """Get a new integer ID."""
      while True:
          new_id = random.randint(0, sys.maxsize)
          if new_id not in d:
              break
      return new_id
  
  
  @contextmanager
  def temp_set_attrs(obj: Any, **kwargs: Any) -> Generator:
      """Temporary setter.
  
      Utility class for setting a temporary value for an attribute on a class.
      Taken from: https://tinyurl.com/2p89xymh
  
      """
      prev_values = {k: getattr(obj, k) for k in kwargs}
      for k, v in kwargs.items():
          setattr(obj, k, v)
      try:
          yield
      finally:
          for k, v in prev_values.items():
              setattr(obj, k, v)
  
  
  @dataclass
  class ErrorToRetry:
      """Exception types that should be retried.
  
      Args:
          exception_cls (Type[Exception]): Class of exception.
          check_fn (Optional[Callable[[Any]], bool]]):
              A function that takes an exception instance as input and returns
              whether to retry.
  
      """
  
      exception_cls: Type[Exception]
      check_fn: Optional[Callable[[Any], bool]] = None
  
  
  def retry_on_exceptions_with_backoff(
      lambda_fn: Callable,
      errors_to_retry: List[ErrorToRetry],
      max_tries: int = 10,
      min_backoff_secs: float = 0.5,
      max_backoff_secs: float = 60.0,
  ) -> Any:
      """Execute lambda function with retries and exponential backoff.
  
      Args:
          lambda_fn (Callable): Function to be called and output we want.
          errors_to_retry (List[ErrorToRetry]): List of errors to retry.
              At least one needs to be provided.
          max_tries (int): Maximum number of tries, including the first. Defaults to 10.
          min_backoff_secs (float): Minimum amount of backoff time between attempts.
              Defaults to 0.5.
          max_backoff_secs (float): Maximum amount of backoff time between attempts.
              Defaults to 60.
  
      """
      if not errors_to_retry:
          raise ValueError("At least one error to retry needs to be provided")
  
      error_checks = {
          error_to_retry.exception_cls: error_to_retry.check_fn
          for error_to_retry in errors_to_retry
      }
      exception_class_tuples = tuple(error_checks.keys())
  
      backoff_secs = min_backoff_secs
      tries = 0
  
      while True:
          try:
              return lambda_fn()
          except exception_class_tuples as e:
              traceback.print_exc()
              tries += 1
              if tries >= max_tries:
                  raise
              check_fn = error_checks.get(e.__class__)
              if check_fn and not check_fn(e):
                  raise
              time.sleep(backoff_secs)
              backoff_secs = min(backoff_secs * 2, max_backoff_secs)
  
  
  def truncate_text(text: str, max_length: int) -> str:
      """Truncate text to a maximum length."""
      if len(text) <= max_length:
          return text
      return text[: max_length - 3] + "..."
  
  
  def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
      """Iterate over an iterable in batches.
  
      >>> list(iter_batch([1,2,3,4,5], 3))
      [[1, 2, 3], [4, 5]]
      """
      source_iter = iter(iterable)
      while source_iter:
          b = list(islice(source_iter, size))
          if len(b) == 0:
              break
          yield b
  
  
  def concat_dirs(dirname: str, basename: str) -> str:
      """
      Append basename to dirname, avoiding backslashes when running on windows.
  
      os.path.join(dirname, basename) will add a backslash before dirname if
      basename does not end with a slash, so we make sure it does.
      """
      dirname += "/" if dirname[-1] != "/" else ""
      return os.path.join(dirname, basename)
  
  
  def get_tqdm_iterable(items: Iterable, show_progress: bool, desc: str) -> Iterable:
      """
      Optionally get a tqdm iterable. Ensures tqdm.auto is used.
      """
      _iterator = items
      if show_progress:
          try:
              from tqdm.auto import tqdm
  
              return tqdm(items, desc=desc)
          except ImportError:
              pass
      return _iterator
  
  
  def count_tokens(text: str) -> int:
      tokenizer = get_tokenizer()
      tokens = tokenizer(text)
      return len(tokens)
  
  
  def get_transformer_tokenizer_fn(model_name: str) -> Callable[[str], List[str]]:
      """
      Args:
          model_name(str): the model name of the tokenizer.
                          For instance, fxmarty/tiny-llama-fast-tokenizer.
      """
      try:
          from transformers import AutoTokenizer
      except ImportError:
          raise ValueError(
              "`transformers` package not found, please run `pip install transformers`"
          )
      tokenizer = AutoTokenizer.from_pretrained(model_name)
      return tokenizer.tokenize
  
  
  def get_cache_dir() -> str:
      """Locate a platform-appropriate cache directory for llama_index,
      and create it if it doesn't yet exist.
      """
      # User override
      if "LLAMA_INDEX_CACHE_DIR" in os.environ:
          path = Path(os.environ["LLAMA_INDEX_CACHE_DIR"])
  
      # Linux, Unix, AIX, etc.
      elif os.name == "posix" and sys.platform != "darwin":
          path = Path("/tmp/llama_index")
  
      # Mac OS
      elif sys.platform == "darwin":
          path = Path(os.path.expanduser("~"), "Library/Caches/llama_index")
  
      # Windows (hopefully)
      else:
          local = os.environ.get("LOCALAPPDATA", None) or os.path.expanduser(
              "~\\AppData\\Local"
          )
          path = Path(local, "llama_index")
  
      if not os.path.exists(path):
          os.makedirs(
              path, exist_ok=True
          )  # prevents https://github.com/jerryjliu/llama_index/issues/7362
      return str(path)
  
  
  def add_sync_version(func: Any) -> Any:
      """Decorator for adding sync version of an async function. The sync version
      is added as a function attribute to the original function, func.
  
      Args:
          func(Any): the async function for which a sync variant will be built.
      """
      assert asyncio.iscoroutinefunction(func)
  
      @wraps(func)
      def _wrapper(*args: Any, **kwds: Any) -> Any:
          return asyncio.get_event_loop().run_until_complete(func(*args, **kwds))
  
      func.sync = _wrapper
      return func
  
  
  # Sample text from llama_index's readme
  SAMPLE_TEXT = """
  Context
  LLMs are a phenomenal piece of technology for knowledge generation and reasoning.
  They are pre-trained on large amounts of publicly available data.
  How do we best augment LLMs with our own private data?
  We need a comprehensive toolkit to help perform this data augmentation for LLMs.
  
  Proposed Solution
  That's where LlamaIndex comes in. LlamaIndex is a "data framework" to help
  you build LLM  apps. It provides the following tools:
  
  Offers data connectors to ingest your existing data sources and data formats
  (APIs, PDFs, docs, SQL, etc.)
  Provides ways to structure your data (indices, graphs) so that this data can be
  easily used with LLMs.
  Provides an advanced retrieval/query interface over your data:
  Feed in any LLM input prompt, get back retrieved context and knowledge-augmented output.
  Allows easy integrations with your outer application framework
  (e.g. with LangChain, Flask, Docker, ChatGPT, anything else).
  LlamaIndex provides tools for both beginner users and advanced users.
  Our high-level API allows beginner users to use LlamaIndex to ingest and
  query their data in 5 lines of code. Our lower-level APIs allow advanced users to
  customize and extend any module (data connectors, indices, retrievers, query engines,
  reranking modules), to fit their needs.
  """
  
  _LLAMA_INDEX_COLORS = {
      "llama_pink": "38;2;237;90;200",
      "llama_blue": "38;2;90;149;237",
      "llama_turquoise": "38;2;11;159;203",
      "llama_lavender": "38;2;155;135;227",
  }
  
  _ANSI_COLORS = {
      "red": "31",
      "green": "32",
      "yellow": "33",
      "blue": "34",
      "magenta": "35",
      "cyan": "36",
      "pink": "38;5;200",
  }
  
  
  def get_color_mapping(
      items: List[str], use_llama_index_colors: bool = True
  ) -> Dict[str, str]:
      """
      Get a mapping of items to colors.
  
      Args:
          items (List[str]): List of items to be mapped to colors.
          use_llama_index_colors (bool, optional): Flag to indicate
          whether to use LlamaIndex colors or ANSI colors.
              Defaults to True.
  
      Returns:
          Dict[str, str]: Mapping of items to colors.
      """
      if use_llama_index_colors:
          color_palette = _LLAMA_INDEX_COLORS
      else:
          color_palette = _ANSI_COLORS
  
      colors = list(color_palette.keys())
      return {item: colors[i % len(colors)] for i, item in enumerate(items)}
  
  
  def _get_colored_text(text: str, color: str) -> str:
      """
      Get the colored version of the input text.
  
      Args:
          text (str): Input text.
          color (str): Color to be applied to the text.
  
      Returns:
          str: Colored version of the input text.
      """
      all_colors = {**_LLAMA_INDEX_COLORS, **_ANSI_COLORS}
  
      if color not in all_colors:
          return f"\033[1;3m{text}\033[0m"  # just bolded and italicized
  
      color = all_colors[color]
  
      return f"\033[1;3;{color}m{text}\033[0m"
  
  
  def print_text(text: str, color: Optional[str] = None, end: str = "") -> None:
      """
      Print the text with the specified color.
  
      Args:
          text (str): Text to be printed.
          color (str, optional): Color to be applied to the text. Supported colors are:
              llama_pink, llama_blue, llama_turquoise, llama_lavender,
              red, green, yellow, blue, magenta, cyan, pink.
          end (str, optional): String appended after the last character of the text.
  
      Returns:
          None
      """
      text_to_print = _get_colored_text(text, color) if color is not None else text
      print(text_to_print, end=end)
  
  
  def infer_torch_device() -> str:
      """Infer the input to torch.device."""
      try:
          has_cuda = torch.cuda.is_available()
      except NameError:
          import torch
  
          has_cuda = torch.cuda.is_available()
      if has_cuda:
          return "cuda"
      if torch.backends.mps.is_available():
          return "mps"
      return "cpu"
  
  
  def unit_generator(x: Any) -> Generator[Any, None, None]:
      """A function that returns a generator of a single element.
  
      Args:
          x (Any): the element to build yield
  
      Yields:
          Any: the single element
      """
      yield x
  
  
  async def async_unit_generator(x: Any) -> AsyncGenerator[Any, None]:
      """A function that returns a generator of a single element.
  
      Args:
          x (Any): the element to build yield
  
      Yields:
          Any: the single element
      """
      yield x
  
apis: llama_index.utils.get_cache_dir
extract_api: [((7192, 7223), 'os.path.join', 'os.path.join', (['dirname', 'basename'], {}), '(dirname, basename)\n', (7204, 7223), False, 'import os\n'), ((8174, 8215), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['model_name'], {}), '(model_name)\n', (8203, 8215), False, 'from transformers import AutoTokenizer\n'), ((9451, 9484), 'asyncio.iscoroutinefunction', 'asyncio.iscoroutinefunction', (['func'], {}), '(func)\n', (9478, 9484), False, 'import asyncio\n'), ((9491, 9502), 'functools.wraps', 'wraps', (['func'], {}), '(func)\n', (9496, 9502), False, 'from functools import partial, wraps\n'), ((13466, 13499), 'torch.backends.mps.is_available', 'torch.backends.mps.is_available', ([], {}), '()\n', (13497, 13499), False, 'import torch\n'), ((3225, 3269), 'tiktoken.encoding_for_model', 'tiktoken.encoding_for_model', (['"""gpt-3.5-turbo"""'], {}), "('gpt-3.5-turbo')\n", (3252, 3269), False, 'import tiktoken\n'), ((3290, 3332), 'functools.partial', 'partial', (['enc.encode'], {'allowed_special': '"""all"""'}), "(enc.encode, allowed_special='all')\n", (3297, 3332), False, 'from functools import partial, wraps\n'), ((3740, 3770), 'random.randint', 'random.randint', (['(0)', 'sys.maxsize'], {}), '(0, sys.maxsize)\n', (3754, 3770), False, 'import random\n'), ((8478, 8519), 'pathlib.Path', 'Path', (["os.environ['LLAMA_INDEX_CACHE_DIR']"], {}), "(os.environ['LLAMA_INDEX_CACHE_DIR'])\n", (8482, 8519), False, 'from pathlib import Path\n'), ((8981, 9001), 'os.path.exists', 'os.path.exists', (['path'], {}), '(path)\n', (8995, 9001), False, 'import os\n'), ((9011, 9043), 'os.makedirs', 'os.makedirs', (['path'], {'exist_ok': '(True)'}), '(path, exist_ok=True)\n', (9022, 9043), False, 'import os\n'), ((13305, 13330), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (13328, 13330), False, 'import torch\n'), ((1200, 1229), 'tiktoken.get_encoding', 'tiktoken.get_encoding', (['"""gpt2"""'], {}), "('gpt2')\n", (1221, 1229), False, 'import tiktoken\n'), ((1260, 1299), 'typing.cast', 'cast', (['Callable[[str], List]', 'enc.encode'], {}), '(Callable[[str], List], enc.encode)\n', (1264, 1299), False, 'from typing import Any, AsyncGenerator, Callable, Dict, Generator, Iterable, List, Optional, Protocol, Set, Type, Union, cast, runtime_checkable\n'), ((1330, 1377), 'functools.partial', 'partial', (['self._tokenizer'], {'allowed_special': '"""all"""'}), "(self._tokenizer, allowed_special='all')\n", (1337, 1377), False, 'from functools import partial, wraps\n'), ((1883, 1898), 'llama_index.utils.get_cache_dir', 'get_cache_dir', ([], {}), '()\n', (1896, 1898), False, 'from llama_index.utils import get_cache_dir\n'), ((1927, 1965), 'os.environ.get', 'os.environ.get', (['"""NLTK_DATA"""', 'cache_dir'], {}), "('NLTK_DATA', cache_dir)\n", (1941, 1965), False, 'import os\n'), ((2341, 2367), 'nltk.corpus.stopwords.words', 'stopwords.words', (['"""english"""'], {}), "('english')\n", (2356, 2367), False, 'from nltk.corpus import stopwords\n'), ((3560, 3572), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (3570, 3572), False, 'import uuid\n'), ((6757, 6782), 'itertools.islice', 'islice', (['source_iter', 'size'], {}), '(source_iter, size)\n', (6763, 6782), False, 'from itertools import islice\n'), ((7505, 7527), 'tqdm.auto.tqdm', 'tqdm', (['items'], {'desc': 'desc'}), '(items, desc=desc)\n', (7509, 7527), False, 'from tqdm.auto import tqdm\n'), ((8623, 8647), 'pathlib.Path', 'Path', (['"""/tmp/llama_index"""'], {}), "('/tmp/llama_index')\n", (8627, 8647), False, 'from pathlib import Path\n'), ((13394, 13419), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (13417, 13419), False, 'import torch\n'), ((2101, 2137), 'nltk.data.path.append', 'nltk.data.path.append', (['nltk_data_dir'], {}), '(nltk_data_dir)\n', (2122, 2137), False, 'import nltk\n'), ((2172, 2207), 'nltk.data.find', 'nltk.data.find', (['"""corpora/stopwords"""'], {}), "('corpora/stopwords')\n", (2186, 2207), False, 'import nltk\n'), ((5966, 5987), 'traceback.print_exc', 'traceback.print_exc', ([], {}), '()\n', (5985, 5987), False, 'import traceback\n'), ((6200, 6224), 'time.sleep', 'time.sleep', (['backoff_secs'], {}), '(backoff_secs)\n', (6210, 6224), False, 'import time\n'), ((8942, 8968), 'pathlib.Path', 'Path', (['local', '"""llama_index"""'], {}), "(local, 'llama_index')\n", (8946, 8968), False, 'from pathlib import Path\n'), ((9568, 9592), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (9590, 9592), False, 'import asyncio\n'), ((2256, 2310), 'nltk.download', 'nltk.download', (['"""stopwords"""'], {'download_dir': 'nltk_data_dir'}), "('stopwords', download_dir=nltk_data_dir)\n", (2269, 2310), False, 'import nltk\n'), ((8717, 8740), 'os.path.expanduser', 'os.path.expanduser', (['"""~"""'], {}), "('~')\n", (8735, 8740), False, 'import os\n'), ((8825, 8861), 'os.environ.get', 'os.environ.get', (['"""LOCALAPPDATA"""', 'None'], {}), "('LOCALAPPDATA', None)\n", (8839, 8861), False, 'import os\n'), ((8865, 8904), 'os.path.expanduser', 'os.path.expanduser', (['"""~\\\\AppData\\\\Local"""'], {}), "('~\\\\AppData\\\\Local')\n", (8883, 8904), False, 'import os\n')]