code: import lancedb
  import os
  import gradio as gr
  from sentence_transformers import SentenceTransformer
  from transformers import AutoTokenizer, AutoModelForSequenceClassification
  import torch
  
  # For Text Similarity and Relevance Ranking:
  # valhalla/distilbart-mnli-12-3
  # sentence-transformers/cross-encoder/stsb-roberta-large
  #
  # For Question Answering:
  # deepset/roberta-base-squad2
  # cross-encoder/quora-distilroberta-base
  
  CROSS_ENC_MODEL = os.getenv("CROSS_ENC_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2")
  
  # Initialize the tokenizer and model for reranking
  tokenizer = AutoTokenizer.from_pretrained(CROSS_ENC_MODEL)
  cross_encoder = AutoModelForSequenceClassification.from_pretrained(CROSS_ENC_MODEL)
  cross_encoder.eval()  # Put model in evaluation mode
  
  db = lancedb.connect(".lancedb")
  TABLE = db.open_table(os.getenv("TABLE_NAME"))
  VECTOR_COLUMN = os.getenv("VECTOR_COLUMN", "vector")
  TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
  BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
  
  retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
  
  def rerank(query, documents):
      pairs = [[query, doc] for doc in documents]  # Create pairs of query and each document
      inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors="pt")
      with torch.no_grad():
          scores = cross_encoder(**inputs).logits.squeeze()  # Get scores for each pair
      sorted_docs = [doc for _, doc in sorted(zip(scores, documents), key=lambda x: x[0], reverse=True)]
      return sorted_docs
  
  def retrieve(query, k, rr=True):
      query_vec = retriever.encode(query)
      try:
          documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
          documents = [doc[TEXT_COLUMN] for doc in documents]
  
          # Rerank the retrieved documents if rr (rerank) is True
          if rr:
              documents = rerank(query, documents)
  
          return documents
  
      except Exception as e:
          raise gr.Error(str(e))
  
apis: lancedb.connect
extract_api: [((440, 508), 'os.getenv', 'os.getenv', (['"""CROSS_ENC_MODEL"""', '"""cross-encoder/ms-marco-MiniLM-L-6-v2"""'], {}), "('CROSS_ENC_MODEL', 'cross-encoder/ms-marco-MiniLM-L-6-v2')\n", (449, 508), False, 'import os\n'), ((573, 619), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['CROSS_ENC_MODEL'], {}), '(CROSS_ENC_MODEL)\n', (602, 619), False, 'from transformers import AutoTokenizer, AutoModelForSequenceClassification\n'), ((636, 703), 'transformers.AutoModelForSequenceClassification.from_pretrained', 'AutoModelForSequenceClassification.from_pretrained', (['CROSS_ENC_MODEL'], {}), '(CROSS_ENC_MODEL)\n', (686, 703), False, 'from transformers import AutoTokenizer, AutoModelForSequenceClassification\n'), ((763, 790), 'lancedb.connect', 'lancedb.connect', (['""".lancedb"""'], {}), "('.lancedb')\n", (778, 790), False, 'import lancedb\n'), ((854, 890), 'os.getenv', 'os.getenv', (['"""VECTOR_COLUMN"""', '"""vector"""'], {}), "('VECTOR_COLUMN', 'vector')\n", (863, 890), False, 'import os\n'), ((905, 937), 'os.getenv', 'os.getenv', (['"""TEXT_COLUMN"""', '"""text"""'], {}), "('TEXT_COLUMN', 'text')\n", (914, 937), False, 'import os\n'), ((813, 836), 'os.getenv', 'os.getenv', (['"""TABLE_NAME"""'], {}), "('TABLE_NAME')\n", (822, 836), False, 'import os\n'), ((955, 982), 'os.getenv', 'os.getenv', (['"""BATCH_SIZE"""', '(32)'], {}), "('BATCH_SIZE', 32)\n", (964, 982), False, 'import os\n'), ((1017, 1039), 'os.getenv', 'os.getenv', (['"""EMB_MODEL"""'], {}), "('EMB_MODEL')\n", (1026, 1039), False, 'import os\n'), ((1254, 1269), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (1267, 1269), False, 'import torch\n')]