code: # Wrapper for Hugging Face APIs for llmlib
  
  from llmlib.base_model_wrapper import BaseModelWrapper
  
  from llama_index import ListIndex, SimpleDirectoryReader
  from langchain.embeddings.huggingface import HuggingFaceEmbeddings
  from llama_index import LangchainEmbedding, ServiceContext
  from llama_index import ListIndex, PromptHelper, SimpleDirectoryReader
  from transformers import pipeline
  
  import torch
  from langchain.llms.base import LLM
  from llama_index import LLMPredictor
  
  from pprint import pprint
  
  class CustomLLM(LLM):
      model_name = "facebook/opt-iml-1.3b"
      # I am not using a GPU, but you can add device="cuda:0"
      # to the pipeline call if you have a local GPU or
      # are running this on Google Colab:
      pipeline = pipeline("text-generation",
                          model=model_name,
                          model_kwargs={"torch_dtype":torch.bfloat16})
  
      def _call(self, prompt, stop = None):
          prompt_length = len(prompt)
          response = self.pipeline(prompt, max_new_tokens=200)
          pprint(response)
          first_response = response[0]["generated_text"]
          # only return newly generated tokens
          returned_text = first_response[prompt_length:]
          return returned_text
  
      @property
      def _identifying_params(self):
          return {"name_of_model": self.model_name}
  
      @property
      def _llm_type(self):
          return "custom"
  
  class HuggingFaceAiWrapper(BaseModelWrapper):
      def __init__(self, key=None, embeddings_dir="./db_embeddings"):
          super().__init__(embeddings_dir=embeddings_dir)
          self.llm_predictor = LLMPredictor(llm=CustomLLM())
          self.embed_model = LangchainEmbedding(HuggingFaceEmbeddings())
          self.service_context = \
            ServiceContext.from_defaults(llm_predictor=self.llm_predictor,
                                         embed_model=self.embed_model)
          max_input_size = 512
          num_output = 64
          max_chunk_overlap = 0 # 10
          self.prompt_helper = PromptHelper(max_input_size, num_output,
                                            max_chunk_overlap)
          self.pipeline = None
  
      # complete text:
      def get_completion(self, prompt, max_tokens=64):
          if self.pipeline is None:
              self.pipeline = pipeline("text-generation",
                                       model="facebook/opt-iml-1.3b",
                                       model_kwargs={"torch_dtype":torch.bfloat16})   
          c = self.pipeline(prompt, max_new_tokens=max_tokens)
          pprint(c)
          try:
              return c[0]["generated_text"]
          except Exception as e:
              print(e)
              return ""
          
  
      def create_local_embeddings_files_in_dir(self, path):
          " path is a directory "
          self.documents = SimpleDirectoryReader(path).load_data()
          self.index = ListIndex.from_documents(documents=self.documents, 
                                                llm_predictor=self.llm_predictor,
                                                prompt_helper=self.prompt_helper)
          self.index = self.index.as_query_engine(llm_predictor=self.llm_predictor)
  
      # query local embeddings:
      def query_local_embeddings(self, query, n=10):
          answer = self.index.query(query)
          return answer
apis: langchain.embeddings.huggingface.HuggingFaceEmbeddings
extract_api: [((735, 830), 'transformers.pipeline', 'pipeline', (['"""text-generation"""'], {'model': 'model_name', 'model_kwargs': "{'torch_dtype': torch.bfloat16}"}), "('text-generation', model=model_name, model_kwargs={'torch_dtype':\n    torch.bfloat16})\n", (743, 830), False, 'from transformers import pipeline\n'), ((1022, 1038), 'pprint.pprint', 'pprint', (['response'], {}), '(response)\n', (1028, 1038), False, 'from pprint import pprint\n'), ((1731, 1828), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm_predictor': 'self.llm_predictor', 'embed_model': 'self.embed_model'}), '(llm_predictor=self.llm_predictor, embed_model=\n    self.embed_model)\n', (1759, 1828), False, 'from llama_index import LangchainEmbedding, ServiceContext\n'), ((1980, 2039), 'llama_index.PromptHelper', 'PromptHelper', (['max_input_size', 'num_output', 'max_chunk_overlap'], {}), '(max_input_size, num_output, max_chunk_overlap)\n', (1992, 2039), False, 'from llama_index import ListIndex, PromptHelper, SimpleDirectoryReader\n'), ((2498, 2507), 'pprint.pprint', 'pprint', (['c'], {}), '(c)\n', (2504, 2507), False, 'from pprint import pprint\n'), ((2823, 2946), 'llama_index.ListIndex.from_documents', 'ListIndex.from_documents', ([], {'documents': 'self.documents', 'llm_predictor': 'self.llm_predictor', 'prompt_helper': 'self.prompt_helper'}), '(documents=self.documents, llm_predictor=self.\n    llm_predictor, prompt_helper=self.prompt_helper)\n', (2847, 2946), False, 'from llama_index import ListIndex, PromptHelper, SimpleDirectoryReader\n'), ((1663, 1686), 'langchain.embeddings.huggingface.HuggingFaceEmbeddings', 'HuggingFaceEmbeddings', ([], {}), '()\n', (1684, 1686), False, 'from langchain.embeddings.huggingface import HuggingFaceEmbeddings\n'), ((2248, 2357), 'transformers.pipeline', 'pipeline', (['"""text-generation"""'], {'model': '"""facebook/opt-iml-1.3b"""', 'model_kwargs': "{'torch_dtype': torch.bfloat16}"}), "('text-generation', model='facebook/opt-iml-1.3b', model_kwargs={\n    'torch_dtype': torch.bfloat16})\n", (2256, 2357), False, 'from transformers import pipeline\n'), ((2762, 2789), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', (['path'], {}), '(path)\n', (2783, 2789), False, 'from llama_index import ListIndex, PromptHelper, SimpleDirectoryReader\n')]