code: #
  # Copyright 2016 The BigDL Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  
  import torch
  from llama_index.embeddings.huggingface import HuggingFaceEmbedding
  from sqlalchemy import make_url
  from llama_index.vector_stores.postgres import PGVectorStore
  # from llama_index.llms.llama_cpp import LlamaCPP
  import psycopg2
  from pathlib import Path
  from llama_index.readers.file import PyMuPDFReader
  from llama_index.core.schema import NodeWithScore
  from typing import Optional
  from llama_index.core.query_engine import RetrieverQueryEngine
  from llama_index.core import QueryBundle
  from llama_index.core.retrievers import BaseRetriever
  from typing import Any, List
  from llama_index.core.node_parser import SentenceSplitter
  from llama_index.core.vector_stores import VectorStoreQuery
  import argparse
  
  def load_vector_database(username, password):
      db_name = "example_db"
      host = "localhost"
      password = password
      port = "5432"
      user = username
      # conn = psycopg2.connect(connection_string)
      conn = psycopg2.connect(
          dbname="postgres",
          host=host,
          password=password,
          port=port,
          user=user,
      )
      conn.autocommit = True
  
      with conn.cursor() as c:
          c.execute(f"DROP DATABASE IF EXISTS {db_name}")
          c.execute(f"CREATE DATABASE {db_name}")
      
      vector_store = PGVectorStore.from_params(
          database=db_name,
          host=host,
          password=password,
          port=port,
          user=user,
          table_name="llama2_paper",
          embed_dim=384,  # openai embedding dimension
      )
      return vector_store
  
  
  def load_data(data_path):
      loader = PyMuPDFReader()
      documents = loader.load(file_path=data_path)
  
  
      text_parser = SentenceSplitter(
          chunk_size=1024,
          # separator=" ",
      )
      text_chunks = []
      # maintain relationship with source doc index, to help inject doc metadata in (3)
      doc_idxs = []
      for doc_idx, doc in enumerate(documents):
          cur_text_chunks = text_parser.split_text(doc.text)
          text_chunks.extend(cur_text_chunks)
          doc_idxs.extend([doc_idx] * len(cur_text_chunks))
  
      from llama_index.core.schema import TextNode
      nodes = []
      for idx, text_chunk in enumerate(text_chunks):
          node = TextNode(
              text=text_chunk,
          )
          src_doc = documents[doc_idxs[idx]]
          node.metadata = src_doc.metadata
          nodes.append(node)
      return nodes
      
  
  
  
  
  
  class VectorDBRetriever(BaseRetriever):
      """Retriever over a postgres vector store."""
  
      def __init__(
          self,
          vector_store: PGVectorStore,
          embed_model: Any,
          query_mode: str = "default",
          similarity_top_k: int = 2,
      ) -> None:
          """Init params."""
          self._vector_store = vector_store
          self._embed_model = embed_model
          self._query_mode = query_mode
          self._similarity_top_k = similarity_top_k
          super().__init__()
  
      def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
          """Retrieve."""
          query_embedding = self._embed_model.get_query_embedding(
              query_bundle.query_str
          )
          vector_store_query = VectorStoreQuery(
              query_embedding=query_embedding,
              similarity_top_k=self._similarity_top_k,
              mode=self._query_mode,
          )
          query_result = self._vector_store.query(vector_store_query)
  
          nodes_with_scores = []
          for index, node in enumerate(query_result.nodes):
              score: Optional[float] = None
              if query_result.similarities is not None:
                  score = query_result.similarities[index]
              nodes_with_scores.append(NodeWithScore(node=node, score=score))
  
          return nodes_with_scores
      
  def completion_to_prompt(completion):
      return f"<|system|>\n</s>\n<|user|>\n{completion}</s>\n<|assistant|>\n"
  
  
  # Transform a list of chat messages into zephyr-specific input
  def messages_to_prompt(messages):
      prompt = ""
      for message in messages:
          if message.role == "system":
              prompt += f"<|system|>\n{message.content}</s>\n"
          elif message.role == "user":
              prompt += f"<|user|>\n{message.content}</s>\n"
          elif message.role == "assistant":
              prompt += f"<|assistant|>\n{message.content}</s>\n"
  
      # ensure we start with a system prompt, insert blank if needed
      if not prompt.startswith("<|system|>\n"):
          prompt = "<|system|>\n</s>\n" + prompt
  
      # add final assistant prompt
      prompt = prompt + "<|assistant|>\n"
  
      return prompt
  
  def main(args):
      embed_model = HuggingFaceEmbedding(model_name=args.embedding_model_path)
      
      # Use custom LLM in BigDL
      from bigdl.llm.llamaindex.llms import BigdlLLM
      llm = BigdlLLM(
          model_name=args.model_path,
          tokenizer_name=args.model_path,
          context_window=512,
          max_new_tokens=args.n_predict,
          generate_kwargs={"temperature": 0.7, "do_sample": False},
          model_kwargs={},
          messages_to_prompt=messages_to_prompt,
          completion_to_prompt=completion_to_prompt,
          device_map="xpu",
      )
      
      vector_store = load_vector_database(username=args.user, password=args.password)
      nodes = load_data(data_path=args.data)
      for node in nodes:
          node_embedding = embed_model.get_text_embedding(
              node.get_content(metadata_mode="all")
          )
          node.embedding = node_embedding
      
      vector_store.add(nodes)
      
      # query_str = "Can you tell me about the key concepts for safety finetuning"
      query_str = "Explain about the training data for Llama 2"
      query_embedding = embed_model.get_query_embedding(query_str)
      # construct vector store query
      
  
      query_mode = "default"
      # query_mode = "sparse"
      # query_mode = "hybrid"
  
      vector_store_query = VectorStoreQuery(
          query_embedding=query_embedding, similarity_top_k=2, mode=query_mode
      )
      # returns a VectorStoreQueryResult
      query_result = vector_store.query(vector_store_query)
      # print("Retrieval Results: ")
      # print(query_result.nodes[0].get_content())
  
  
  
      nodes_with_scores = []
      for index, node in enumerate(query_result.nodes):
          score: Optional[float] = None
          if query_result.similarities is not None:
              score = query_result.similarities[index]
          nodes_with_scores.append(NodeWithScore(node=node, score=score))
      
      retriever = VectorDBRetriever(
          vector_store, embed_model, query_mode="default", similarity_top_k=1
      )
      
      
      query_engine = RetrieverQueryEngine.from_args(retriever, llm=llm)
  
      # query_str = "How does Llama 2 perform compared to other open-source models?"
      query_str = args.question
      response = query_engine.query(query_str)
  
  
      print("------------RESPONSE GENERATION---------------------")
      print(str(response))
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser(description='LlamaIndex BigdlLLM Example')
      parser.add_argument('-m','--model-path', type=str, required=True,
                          help='the path to transformers model')
      parser.add_argument('-q', '--question', type=str, default='How does Llama 2 perform compared to other open-source models?',
                          help='qustion you want to ask.')
      parser.add_argument('-d','--data',type=str, default='./data/llama2.pdf',
                          help="the data used during retrieval")
      parser.add_argument('-u', '--user', type=str, required=True,
                          help="user name in the database postgres")
      parser.add_argument('-p','--password', type=str, required=True,
                          help="the password of the user in the database")
      parser.add_argument('-e','--embedding-model-path',default="BAAI/bge-small-en",
                          help="the path to embedding model path")
      parser.add_argument('-n','--n-predict', type=int, default=32,
                          help='max number of predict tokens')
      args = parser.parse_args()
      
      main(args)
apis: llama_index.core.node_parser.SentenceSplitter llama_index.embeddings.huggingface.HuggingFaceEmbedding llama_index.readers.file.PyMuPDFReader llama_index.core.schema.TextNode llama_index.vector_stores.postgres.PGVectorStore.from_params llama_index.core.query_engine.RetrieverQueryEngine.from_args llama_index.core.schema.NodeWithScore llama_index.core.vector_stores.VectorStoreQuery
extract_api: [((1521, 1612), 'psycopg2.connect', 'psycopg2.connect', ([], {'dbname': '"""postgres"""', 'host': 'host', 'password': 'password', 'port': 'port', 'user': 'user'}), "(dbname='postgres', host=host, password=password, port=port,\n    user=user)\n", (1537, 1612), False, 'import psycopg2\n'), ((1841, 1982), 'llama_index.vector_stores.postgres.PGVectorStore.from_params', 'PGVectorStore.from_params', ([], {'database': 'db_name', 'host': 'host', 'password': 'password', 'port': 'port', 'user': 'user', 'table_name': '"""llama2_paper"""', 'embed_dim': '(384)'}), "(database=db_name, host=host, password=password,\n    port=port, user=user, table_name='llama2_paper', embed_dim=384)\n", (1866, 1982), False, 'from llama_index.vector_stores.postgres import PGVectorStore\n'), ((2137, 2152), 'llama_index.readers.file.PyMuPDFReader', 'PyMuPDFReader', ([], {}), '()\n', (2150, 2152), False, 'from llama_index.readers.file import PyMuPDFReader\n'), ((2222, 2255), 'llama_index.core.node_parser.SentenceSplitter', 'SentenceSplitter', ([], {'chunk_size': '(1024)'}), '(chunk_size=1024)\n', (2238, 2255), False, 'from llama_index.core.node_parser import SentenceSplitter\n'), ((5116, 5174), 'llama_index.embeddings.huggingface.HuggingFaceEmbedding', 'HuggingFaceEmbedding', ([], {'model_name': 'args.embedding_model_path'}), '(model_name=args.embedding_model_path)\n', (5136, 5174), False, 'from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n'), ((5271, 5583), 'bigdl.llm.llamaindex.llms.BigdlLLM', 'BigdlLLM', ([], {'model_name': 'args.model_path', 'tokenizer_name': 'args.model_path', 'context_window': '(512)', 'max_new_tokens': 'args.n_predict', 'generate_kwargs': "{'temperature': 0.7, 'do_sample': False}", 'model_kwargs': '{}', 'messages_to_prompt': 'messages_to_prompt', 'completion_to_prompt': 'completion_to_prompt', 'device_map': '"""xpu"""'}), "(model_name=args.model_path, tokenizer_name=args.model_path,\n    context_window=512, max_new_tokens=args.n_predict, generate_kwargs={\n    'temperature': 0.7, 'do_sample': False}, model_kwargs={},\n    messages_to_prompt=messages_to_prompt, completion_to_prompt=\n    completion_to_prompt, device_map='xpu')\n", (5279, 5583), False, 'from bigdl.llm.llamaindex.llms import BigdlLLM\n'), ((6353, 6444), 'llama_index.core.vector_stores.VectorStoreQuery', 'VectorStoreQuery', ([], {'query_embedding': 'query_embedding', 'similarity_top_k': '(2)', 'mode': 'query_mode'}), '(query_embedding=query_embedding, similarity_top_k=2, mode=\n    query_mode)\n', (6369, 6444), False, 'from llama_index.core.vector_stores import VectorStoreQuery\n'), ((7083, 7133), 'llama_index.core.query_engine.RetrieverQueryEngine.from_args', 'RetrieverQueryEngine.from_args', (['retriever'], {'llm': 'llm'}), '(retriever, llm=llm)\n', (7113, 7133), False, 'from llama_index.core.query_engine import RetrieverQueryEngine\n'), ((7428, 7494), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""LlamaIndex BigdlLLM Example"""'}), "(description='LlamaIndex BigdlLLM Example')\n", (7451, 7494), False, 'import argparse\n'), ((2759, 2784), 'llama_index.core.schema.TextNode', 'TextNode', ([], {'text': 'text_chunk'}), '(text=text_chunk)\n', (2767, 2784), False, 'from llama_index.core.schema import TextNode\n'), ((3682, 3800), 'llama_index.core.vector_stores.VectorStoreQuery', 'VectorStoreQuery', ([], {'query_embedding': 'query_embedding', 'similarity_top_k': 'self._similarity_top_k', 'mode': 'self._query_mode'}), '(query_embedding=query_embedding, similarity_top_k=self.\n    _similarity_top_k, mode=self._query_mode)\n', (3698, 3800), False, 'from llama_index.core.vector_stores import VectorStoreQuery\n'), ((6893, 6930), 'llama_index.core.schema.NodeWithScore', 'NodeWithScore', ([], {'node': 'node', 'score': 'score'}), '(node=node, score=score)\n', (6906, 6930), False, 'from llama_index.core.schema import NodeWithScore\n'), ((4191, 4228), 'llama_index.core.schema.NodeWithScore', 'NodeWithScore', ([], {'node': 'node', 'score': 'score'}), '(node=node, score=score)\n', (4204, 4228), False, 'from llama_index.core.schema import NodeWithScore\n')]