code: #!/usr/bin/env python3
  # Copyright (c) 2023 Steve Castellotti
  # This file is part of Urcuchillay and is released under the MIT License.
  # See LICENSE file in the project root for full license information.
  
  import argparse
  import logging
  import os
  import sys
  
  import config
  import utils
  
  try:
      import llama_index
      import transformers
  except ModuleNotFoundError as e:
      print('\nError importing Python module(s)')
      print('If installed using setup.sh it may be necessary to run:\n')
      print('pyenv activate urcuchillay-env\n')
      sys.exit(1)
  
  
  DEFAULT_PROMPT = 'What is Urcuchillay?'
  
  
  class Query:
      def __init__(self, args):
  
          self.debug = args.debug
  
          data_path = args.data
          model_path = args.path
          storage_path = args.storage
  
          load = args.load
          save = args.save
          enable_gpu = args.cpu
          temperature = args.temperature
          max_new_tokens = args.max_new_tokens
          context_window = args.context
  
          self.model_name = args.model
          model_url = args.model_url
          embed_model_name = args.embed_model_name
          embed_model_provider = args.embed_model_provider
          pretrained_model_name = args.pretrained_model_name
          pretrained_model_provider = args.pretrained_model_provider
  
          level = logging.DEBUG if self.debug else logging.INFO
          logging.basicConfig(stream=sys.stdout, level=level)
          logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
  
          llama_debug = llama_index.callbacks.LlamaDebugHandler(print_trace_on_end=self.debug)
          self.callback_manager = llama_index.callbacks.CallbackManager([llama_debug])
  
          # llama_index will automatically assume models are cached in a subdirectory of the current path named
          # "models" so we need to handle if a user explicitly included "models" at the end of --model_path
          cache_directory = model_path
          if os.path.basename(model_path) == 'models':
              path = os.path.join(model_path, self.model_name)
              cache_directory = os.path.dirname(model_path)
          else:
              path = os.path.join(model_path, 'models', self.model_name)
          if os.path.exists(path):
              model_url = 'file://' + str(path)
  
          os.environ['LLAMA_INDEX_CACHE_DIR'] = cache_directory
  
          if pretrained_model_name is not None:
              llama_index.set_global_tokenizer(
                  transformers.AutoTokenizer.from_pretrained(
                      pretrained_model_provider + '/' + pretrained_model_name
                  ).encode
              )
  
          self.llm = llama_index.llms.LlamaCPP(
              model_url=model_url,
              model_path=None,  # Note: setting a custom model_path here causes a fault
              temperature=temperature,
              max_new_tokens=max_new_tokens,
              context_window=context_window,
              generate_kwargs={},  # kwargs to pass to __call__()
              model_kwargs={'n_gpu_layers': enable_gpu},  # kwargs to pass to __init__()
              # transform inputs into Llama 2 format
              messages_to_prompt=llama_index.llms.llama_utils.messages_to_prompt,
              completion_to_prompt=llama_index.llms.llama_utils.completion_to_prompt,
              verbose=self.debug,
          )
  
          if args.embed_model_name == 'local':
              embed_model = args.embed_model_name
          else:
              # use Huggingface embeddings
              embed_model = llama_index.embeddings.HuggingFaceEmbedding(
                  model_name=embed_model_provider + '/' + embed_model_name)
  
          # create a service context
          service_context = llama_index.ServiceContext.from_defaults(
              llm=self.llm,
              embed_model=embed_model,
              callback_manager=self.callback_manager
          )
  
          if load:
              # load vector index from storage
              storage_context = llama_index.StorageContext.from_defaults(persist_dir=storage_path)
              index = llama_index.load_index_from_storage(storage_context, service_context=service_context)
          else:
              # load documents
              documents = llama_index.SimpleDirectoryReader(data_path).load_data()
              # create vector store index
              index = llama_index.VectorStoreIndex.from_documents(
                  documents, service_context=service_context
              )
  
          # persist the index to disk
          if save:
              index.storage_context.persist(persist_dir=storage_path)
  
          # set up query engine
          self.query_engine = index.as_query_engine()
  
      def display_exchange(self, query):
          print('Query: %s\n' % query)
  
          if self.model_name in config.Models.MODELS.keys():
              query = config.Models.MODELS[self.model_name]['prompt_template'].replace('{prompt}', query)
              if self.debug:
                  print('Query (prompt): %s\n' % query)
  
          response = self.query_engine.query(query)
          print('Response: %s\n' % str(response).strip())
  
  
  def parse_arguments():
      parser = argparse.ArgumentParser(description='Process command parameters')
      parser = utils.parse_arguments_common(parser)
  
      parser.add_argument('-p', '--prompt', type=str, default=DEFAULT_PROMPT,
                          help='The prompt to process (default: %(default)s)')
      parser.add_argument('--embed_model_name', type=str, default=config.Config.EMBED_MODEL_NAME,
                          help='The name of the embedding model to use (default: %(default)s)')
      parser.add_argument('--embed_model_provider', type=str, default=None,
                          help='The provider of the embedding model to use (default: %(default)s)')
      parser.add_argument('--pretrained_model_name', type=str, default=None,
                          help='The name of the pretrained model to use (default: %(default)s)')
      parser.add_argument('--pretrained_model_provider', type=str, default=None,
                          help='The provider of the pretrained model to use (default: %(default)s)')
  
      args = parser.parse_args()
      args = utils.update_arguments_common(args)
      return args
  
  
  def main():
      args = parse_arguments()
      llm_query = Query(args=args)
      llm_query.display_exchange(args.prompt)
  
  
  if __name__ == '__main__':
      main()
  
apis: llama_index.VectorStoreIndex.from_documents llama_index.callbacks.LlamaDebugHandler llama_index.ServiceContext.from_defaults llama_index.StorageContext.from_defaults llama_index.SimpleDirectoryReader llama_index.embeddings.HuggingFaceEmbedding llama_index.llms.LlamaCPP llama_index.callbacks.CallbackManager llama_index.load_index_from_storage
extract_api: [((5042, 5107), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Process command parameters"""'}), "(description='Process command parameters')\n", (5065, 5107), False, 'import argparse\n'), ((5121, 5157), 'utils.parse_arguments_common', 'utils.parse_arguments_common', (['parser'], {}), '(parser)\n', (5149, 5157), False, 'import utils\n'), ((6065, 6100), 'utils.update_arguments_common', 'utils.update_arguments_common', (['args'], {}), '(args)\n', (6094, 6100), False, 'import utils\n'), ((541, 552), 'sys.exit', 'sys.exit', (['(1)'], {}), '(1)\n', (549, 552), False, 'import sys\n'), ((1350, 1401), 'logging.basicConfig', 'logging.basicConfig', ([], {'stream': 'sys.stdout', 'level': 'level'}), '(stream=sys.stdout, level=level)\n', (1369, 1401), False, 'import logging\n'), ((1506, 1576), 'llama_index.callbacks.LlamaDebugHandler', 'llama_index.callbacks.LlamaDebugHandler', ([], {'print_trace_on_end': 'self.debug'}), '(print_trace_on_end=self.debug)\n', (1545, 1576), False, 'import llama_index\n'), ((1609, 1661), 'llama_index.callbacks.CallbackManager', 'llama_index.callbacks.CallbackManager', (['[llama_debug]'], {}), '([llama_debug])\n', (1646, 1661), False, 'import llama_index\n'), ((2184, 2204), 'os.path.exists', 'os.path.exists', (['path'], {}), '(path)\n', (2198, 2204), False, 'import os\n'), ((2603, 3000), 'llama_index.llms.LlamaCPP', 'llama_index.llms.LlamaCPP', ([], {'model_url': 'model_url', 'model_path': 'None', 'temperature': 'temperature', 'max_new_tokens': 'max_new_tokens', 'context_window': 'context_window', 'generate_kwargs': '{}', 'model_kwargs': "{'n_gpu_layers': enable_gpu}", 'messages_to_prompt': 'llama_index.llms.llama_utils.messages_to_prompt', 'completion_to_prompt': 'llama_index.llms.llama_utils.completion_to_prompt', 'verbose': 'self.debug'}), "(model_url=model_url, model_path=None, temperature\n    =temperature, max_new_tokens=max_new_tokens, context_window=\n    context_window, generate_kwargs={}, model_kwargs={'n_gpu_layers':\n    enable_gpu}, messages_to_prompt=llama_index.llms.llama_utils.\n    messages_to_prompt, completion_to_prompt=llama_index.llms.llama_utils.\n    completion_to_prompt, verbose=self.debug)\n", (2628, 3000), False, 'import llama_index\n'), ((3636, 3760), 'llama_index.ServiceContext.from_defaults', 'llama_index.ServiceContext.from_defaults', ([], {'llm': 'self.llm', 'embed_model': 'embed_model', 'callback_manager': 'self.callback_manager'}), '(llm=self.llm, embed_model=\n    embed_model, callback_manager=self.callback_manager)\n', (3676, 3760), False, 'import llama_index\n'), ((1441, 1481), 'logging.StreamHandler', 'logging.StreamHandler', ([], {'stream': 'sys.stdout'}), '(stream=sys.stdout)\n', (1462, 1481), False, 'import logging\n'), ((1927, 1955), 'os.path.basename', 'os.path.basename', (['model_path'], {}), '(model_path)\n', (1943, 1955), False, 'import os\n'), ((1988, 2029), 'os.path.join', 'os.path.join', (['model_path', 'self.model_name'], {}), '(model_path, self.model_name)\n', (2000, 2029), False, 'import os\n'), ((2060, 2087), 'os.path.dirname', 'os.path.dirname', (['model_path'], {}), '(model_path)\n', (2075, 2087), False, 'import os\n'), ((2121, 2172), 'os.path.join', 'os.path.join', (['model_path', '"""models"""', 'self.model_name'], {}), "(model_path, 'models', self.model_name)\n", (2133, 2172), False, 'import os\n'), ((3455, 3560), 'llama_index.embeddings.HuggingFaceEmbedding', 'llama_index.embeddings.HuggingFaceEmbedding', ([], {'model_name': "(embed_model_provider + '/' + embed_model_name)"}), "(model_name=embed_model_provider +\n    '/' + embed_model_name)\n", (3498, 3560), False, 'import llama_index\n'), ((3895, 3961), 'llama_index.StorageContext.from_defaults', 'llama_index.StorageContext.from_defaults', ([], {'persist_dir': 'storage_path'}), '(persist_dir=storage_path)\n', (3935, 3961), False, 'import llama_index\n'), ((3982, 4072), 'llama_index.load_index_from_storage', 'llama_index.load_index_from_storage', (['storage_context'], {'service_context': 'service_context'}), '(storage_context, service_context=\n    service_context)\n', (4017, 4072), False, 'import llama_index\n'), ((4252, 4344), 'llama_index.VectorStoreIndex.from_documents', 'llama_index.VectorStoreIndex.from_documents', (['documents'], {'service_context': 'service_context'}), '(documents, service_context=\n    service_context)\n', (4295, 4344), False, 'import llama_index\n'), ((4683, 4710), 'config.Models.MODELS.keys', 'config.Models.MODELS.keys', ([], {}), '()\n', (4708, 4710), False, 'import config\n'), ((1410, 1429), 'logging.getLogger', 'logging.getLogger', ([], {}), '()\n', (1427, 1429), False, 'import logging\n'), ((2424, 2527), 'transformers.AutoTokenizer.from_pretrained', 'transformers.AutoTokenizer.from_pretrained', (["(pretrained_model_provider + '/' + pretrained_model_name)"], {}), "(pretrained_model_provider + '/' +\n    pretrained_model_name)\n", (2466, 2527), False, 'import transformers\n'), ((4135, 4179), 'llama_index.SimpleDirectoryReader', 'llama_index.SimpleDirectoryReader', (['data_path'], {}), '(data_path)\n', (4168, 4179), False, 'import llama_index\n')]