code: import logging
  import json
  import gradio as gr
  import numpy as np
  import lancedb
  import os
  from huggingface_hub import AsyncInferenceClient
  
  
  # Setting up the logging
  logging.basicConfig(level=logging.INFO)
  logger = logging.getLogger(__name__)
  
  # db
  TABLE_NAME = "docs"
  TEXT_COLUMN = "text"
  BATCH_SIZE = int(os.getenv("BATCH_SIZE"))
  NPROBES = int(os.getenv("NPROBES"))
  REFINE_FACTOR = int(os.getenv("REFINE_FACTOR"))
  
  retriever = AsyncInferenceClient(model=os.getenv("EMBED_URL") + "/embed")
  reranker = AsyncInferenceClient(model=os.getenv("RERANK_URL") + "/rerank")
  
  db = lancedb.connect("/usr/src/.lancedb")
  tbl = db.open_table(TABLE_NAME)
  
  
  async def retrieve(query: str, k: int) -> list[str]:
      """
      Retrieve top k items with RETRIEVER
      """
      resp = await retriever.post(
          json={
              "inputs": query,
              "truncate": True
          }
      )
      try:
          query_vec = json.loads(resp)[0]
      except:
          raise gr.Error(resp.decode())
      
      documents = tbl.search(
          query=query_vec
      ).nprobes(NPROBES).refine_factor(REFINE_FACTOR).limit(k).to_list()
      documents = [doc[TEXT_COLUMN] for doc in documents]
  
      return documents
  
  
  async def rerank(query: str, documents: list[str], k: int) -> list[str]:
      """
      Rerank items returned by RETRIEVER and return top k
      """
      scores = []
      for i in range(int(np.ceil(len(documents) / BATCH_SIZE))):
          resp = await reranker.post(
              json={
                  "query": query,
                  "texts": documents[i * BATCH_SIZE:(i + 1) * BATCH_SIZE],
                  "truncate": True
              }
          )
          try:
              batch_scores = json.loads(resp)
              batch_scores = [s["score"] for s in batch_scores]
              scores.extend(batch_scores)
          except:
              raise gr.Error(resp.decode())
      documents = [doc for _, doc in sorted(zip(scores, documents))[-k:]]
  
      return documents
apis: lancedb.connect
extract_api: [((167, 206), 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.INFO'}), '(level=logging.INFO)\n', (186, 206), False, 'import logging\n'), ((216, 243), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (233, 243), False, 'import logging\n'), ((573, 609), 'lancedb.connect', 'lancedb.connect', (['"""/usr/src/.lancedb"""'], {}), "('/usr/src/.lancedb')\n", (588, 609), False, 'import lancedb\n'), ((308, 331), 'os.getenv', 'os.getenv', (['"""BATCH_SIZE"""'], {}), "('BATCH_SIZE')\n", (317, 331), False, 'import os\n'), ((347, 367), 'os.getenv', 'os.getenv', (['"""NPROBES"""'], {}), "('NPROBES')\n", (356, 367), False, 'import os\n'), ((389, 415), 'os.getenv', 'os.getenv', (['"""REFINE_FACTOR"""'], {}), "('REFINE_FACTOR')\n", (398, 415), False, 'import os\n'), ((457, 479), 'os.getenv', 'os.getenv', (['"""EMBED_URL"""'], {}), "('EMBED_URL')\n", (466, 479), False, 'import os\n'), ((530, 553), 'os.getenv', 'os.getenv', (['"""RERANK_URL"""'], {}), "('RERANK_URL')\n", (539, 553), False, 'import os\n'), ((904, 920), 'json.loads', 'json.loads', (['resp'], {}), '(resp)\n', (914, 920), False, 'import json\n'), ((1663, 1679), 'json.loads', 'json.loads', (['resp'], {}), '(resp)\n', (1673, 1679), False, 'import json\n')]