code: import argparse
  import lancedb
  import torch
  import pyarrow as pa
  import pandas as pd
  from pathlib import Path
  import tqdm
  import numpy as np
  import logging
  
  from transformers import AutoConfig
  from sentence_transformers import SentenceTransformer
  
  logging.basicConfig(level=logging.INFO)
  logger = logging.getLogger(__name__)
  
  
  def main():
      parser = argparse.ArgumentParser()
      parser.add_argument("--emb-model", help="embedding model name on HF hub", type=str)
      parser.add_argument("--table", help="table name in DB", type=str)
      parser.add_argument("--input-dir", help="input directory with documents to ingest", type=str)
      parser.add_argument("--vec-column", help="vector column name in the table", type=str, default="vector")
      parser.add_argument("--text-column", help="text column name in the table", type=str, default="text")
      parser.add_argument("--db-loc", help="database location", type=str,
                          default=str(Path().resolve() / ".lancedb"))
      parser.add_argument("--batch-size", help="batch size for embedding model", type=int, default=32)
      parser.add_argument("--num-partitions", help="number of partitions for index", type=int, default=256)
      parser.add_argument("--num-sub-vectors", help="number of sub-vectors for index", type=int, default=96)
  
      args = parser.parse_args()
  
      emb_config = AutoConfig.from_pretrained(args.emb_model)
      emb_dimension = emb_config.hidden_size
  
      assert emb_dimension % args.num_sub_vectors == 0, \
          "Embedding size must be divisible by the num of sub vectors"
  
      model = SentenceTransformer(args.emb_model)
      model.eval()
  
      if torch.backends.mps.is_available():
          device = "mps"
      elif torch.cuda.is_available():
          device = "cuda"
      else:
          device = "cpu"
      logger.info(f"using {str(device)} device")
  
      db = lancedb.connect(args.db_loc)
  
      schema = pa.schema(
        [
            pa.field(args.vec_column, pa.list_(pa.float32(), emb_dimension)),
            pa.field(args.text_column, pa.string())
        ]
      )
      tbl = db.create_table(args.table, schema=schema, mode="overwrite")
  
      input_dir = Path(args.input_dir)
  
      files = list(input_dir.rglob("*"))
      sentences = []
      for file in files:
          if file.is_file():
              with open(file, encoding='utf-8') as f:
                  sentences.append(f.read())
  
      for i in tqdm.tqdm(range(0, int(np.ceil(len(sentences) / args.batch_size)))):
          try:
              batch = [sent for sent in sentences[i * args.batch_size:(i + 1) * args.batch_size] if len(sent) > 0]
              encoded = model.encode(batch, normalize_embeddings=True, device=device)
              encoded = [list(vec) for vec in encoded]
  
              df = pd.DataFrame({
                  args.vec_column: encoded,
                  args.text_column: batch
              })
  
              tbl.add(df)
          except:
              logger.info(f"batch {i} was skipped")
  
      '''
      create ivf-pd index https://lancedb.github.io/lancedb/ann_indexes/
      with the size of the transformer docs, index is not really needed
      but we'll do it for demonstrational purposes
      '''
      tbl.create_index(
          num_partitions=args.num_partitions,
          num_sub_vectors=args.num_sub_vectors,
          vector_column_name=args.vec_column
      )
  
  
  if __name__ == "__main__":
      main()
apis: lancedb.connect
extract_api: [((248, 287), 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.INFO'}), '(level=logging.INFO)\n', (267, 287), False, 'import logging\n'), ((297, 324), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (314, 324), False, 'import logging\n'), ((352, 377), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (375, 377), False, 'import argparse\n'), ((1351, 1393), 'transformers.AutoConfig.from_pretrained', 'AutoConfig.from_pretrained', (['args.emb_model'], {}), '(args.emb_model)\n', (1377, 1393), False, 'from transformers import AutoConfig\n'), ((1576, 1611), 'sentence_transformers.SentenceTransformer', 'SentenceTransformer', (['args.emb_model'], {}), '(args.emb_model)\n', (1595, 1611), False, 'from sentence_transformers import SentenceTransformer\n'), ((1637, 1670), 'torch.backends.mps.is_available', 'torch.backends.mps.is_available', ([], {}), '()\n', (1668, 1670), False, 'import torch\n'), ((1845, 1873), 'lancedb.connect', 'lancedb.connect', (['args.db_loc'], {}), '(args.db_loc)\n', (1860, 1873), False, 'import lancedb\n'), ((2135, 2155), 'pathlib.Path', 'Path', (['args.input_dir'], {}), '(args.input_dir)\n', (2139, 2155), False, 'from pathlib import Path\n'), ((1704, 1729), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (1727, 1729), False, 'import torch\n'), ((2724, 2789), 'pandas.DataFrame', 'pd.DataFrame', (['{args.vec_column: encoded, args.text_column: batch}'], {}), '({args.vec_column: encoded, args.text_column: batch})\n', (2736, 2789), True, 'import pandas as pd\n'), ((2020, 2031), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (2029, 2031), True, 'import pyarrow as pa\n'), ((1952, 1964), 'pyarrow.float32', 'pa.float32', ([], {}), '()\n', (1962, 1964), True, 'import pyarrow as pa\n'), ((955, 961), 'pathlib.Path', 'Path', ([], {}), '()\n', (959, 961), False, 'from pathlib import Path\n')]