code: import argparse
  import os
  import sys
  from concurrent.futures import ProcessPoolExecutor, as_completed
  from functools import lru_cache
  from pathlib import Path
  from typing import Any, Iterator
  
  import lancedb
  import pandas as pd
  import srsly
  from codetiming import Timer
  from dotenv import load_dotenv
  from lancedb.pydantic import pydantic_to_schema
  from sentence_transformers import SentenceTransformer
  from tqdm import tqdm
  
  sys.path.insert(1, os.path.realpath(Path(__file__).resolve().parents[1]))
  from api.config import Settings
  from schemas.wine import LanceModelWine, Wine
  
  
  load_dotenv()
  # Custom types
  JsonBlob = dict[str, Any]
  
  
  class FileNotFoundError(Exception):
      pass
  
  
  @lru_cache()
  def get_settings():
      # Use lru_cache to avoid loading .env file for every request
      return Settings()
  
  
  def chunk_iterable(item_list: list[JsonBlob], chunksize: int) -> Iterator[list[JsonBlob]]:
      """
      Break a large iterable into an iterable of smaller iterables of size `chunksize`
      """
      for i in range(0, len(item_list), chunksize):
          yield item_list[i : i + chunksize]
  
  
  def get_json_data(data_dir: Path, filename: str) -> list[JsonBlob]:
      """Get all line-delimited json files (.jsonl) from a directory with a given prefix"""
      file_path = data_dir / filename
      if not file_path.is_file():
          # File may not have been uncompressed yet so try to do that first
          data = srsly.read_gzip_jsonl(file_path)
          # This time if it isn't there it really doesn't exist
          if not file_path.is_file():
              raise FileNotFoundError(f"No valid .jsonl file found in `{data_dir}`")
      else:
          data = srsly.read_gzip_jsonl(file_path)
      return data
  
  
  def validate(
      data: list[JsonBlob],
      exclude_none: bool = False,
  ) -> list[JsonBlob]:
      validated_data = [Wine(**item).model_dump(exclude_none=exclude_none) for item in data]
      return validated_data
  
  
  def embed_func(batch: list[str], model) -> list[list[float]]:
      return [model.encode(sentence.lower()) for sentence in batch]
  
  
  def vectorize_text(data: list[JsonBlob]) -> list[LanceModelWine] | None:
      # Load a sentence transformer model for semantic similarity from a specified checkpoint
      model_id = get_settings().embedding_model_checkpoint
      assert model_id, "Invalid embedding model checkpoint specified in .env file"
      MODEL = SentenceTransformer(model_id)
  
      ids = [item["id"] for item in data]
      to_vectorize = [text.get("to_vectorize") for text in data]
      vectors = embed_func(to_vectorize, MODEL)
      try:
          data_batch = [{**d, "vector": vector} for d, vector in zip(data, vectors)]
      except Exception as e:
          print(f"{e}: Failed to add ID range {min(ids)}-{max(ids)}")
          return None
      return data_batch
  
  
  def embed_batches(tbl: str, validated_data: list[JsonBlob]) -> pd.DataFrame:
      with ProcessPoolExecutor(max_workers=WORKERS) as executor:
          chunked_data = chunk_iterable(validated_data, CHUNKSIZE)
          embed_data = []
          for chunk in tqdm(chunked_data, total=len(validated_data) // CHUNKSIZE):
              futures = [executor.submit(vectorize_text, chunk)]
              embed_data = [f.result() for f in as_completed(futures) if f.result()][0]
              df = pd.DataFrame.from_dict(embed_data)
              tbl.add(df, mode="overwrite")
  
  
  def main(data: list[JsonBlob]) -> None:
      DB_NAME = f"../{get_settings().lancedb_dir}"
      TABLE = "wines"
      db = lancedb.connect(DB_NAME)
  
      tbl = db.create_table(TABLE, schema=pydantic_to_schema(LanceModelWine), mode="overwrite")
      print(f"Created table `{TABLE}`, with length {len(tbl)}")
  
      with Timer(name="Bulk Index", text="Validated data using Pydantic in {:.4f} sec"):
          validated_data = validate(data, exclude_none=False)
  
      with Timer(name="Embed batches", text="Created sentence embeddings in {:.4f} sec"):
          embed_batches(tbl, validated_data)
  
      print(f"Finished inserting {len(tbl)} items into LanceDB table")
  
      with Timer(name="Create index", text="Created IVF-PQ index in {:.4f} sec"):
          # Creating index (choose num partitions as a power of 2 that's closest to len(dataset) // 5000)
          # In this case, we have 130k datapoints, so the nearest power of 2 is 130000//5000 ~ 32)
          tbl.create_index(metric="cosine", num_partitions=4, num_sub_vectors=32)
  
  
  if __name__ == "__main__":
      # fmt: off
      parser = argparse.ArgumentParser("Bulk index database from the wine reviews JSONL data")
      parser.add_argument("--limit", type=int, default=0, help="Limit the size of the dataset to load for testing purposes")
      parser.add_argument("--chunksize", type=int, default=1000, help="Size of each chunk to break the dataset into before processing")
      parser.add_argument("--filename", type=str, default="winemag-data-130k-v2.jsonl.gz", help="Name of the JSONL zip file to use")
      parser.add_argument("--workers", type=int, default=4, help="Number of workers to use for vectorization")
      args = vars(parser.parse_args())
      # fmt: on
  
      LIMIT = args["limit"]
      DATA_DIR = Path(__file__).parents[3] / "data"
      FILENAME = args["filename"]
      CHUNKSIZE = args["chunksize"]
      WORKERS = args["workers"]
  
      data = list(get_json_data(DATA_DIR, FILENAME))
      assert data, "No data found in the specified file"
      data = data[:LIMIT] if LIMIT > 0 else data
      main(data)
      print("Finished execution!")
  
apis: lancedb.connect lancedb.pydantic.pydantic_to_schema
extract_api: [((580, 593), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (591, 593), False, 'from dotenv import load_dotenv\n'), ((685, 696), 'functools.lru_cache', 'lru_cache', ([], {}), '()\n', (694, 696), False, 'from functools import lru_cache\n'), ((793, 803), 'api.config.Settings', 'Settings', ([], {}), '()\n', (801, 803), False, 'from api.config import Settings\n'), ((2355, 2384), 'sentence_transformers.SentenceTransformer', 'SentenceTransformer', (['model_id'], {}), '(model_id)\n', (2374, 2384), False, 'from sentence_transformers import SentenceTransformer\n'), ((3439, 3463), 'lancedb.connect', 'lancedb.connect', (['DB_NAME'], {}), '(DB_NAME)\n', (3454, 3463), False, 'import lancedb\n'), ((4390, 4469), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (['"""Bulk index database from the wine reviews JSONL data"""'], {}), "('Bulk index database from the wine reviews JSONL data')\n", (4413, 4469), False, 'import argparse\n'), ((1408, 1440), 'srsly.read_gzip_jsonl', 'srsly.read_gzip_jsonl', (['file_path'], {}), '(file_path)\n', (1429, 1440), False, 'import srsly\n'), ((1647, 1679), 'srsly.read_gzip_jsonl', 'srsly.read_gzip_jsonl', (['file_path'], {}), '(file_path)\n', (1668, 1679), False, 'import srsly\n'), ((2852, 2892), 'concurrent.futures.ProcessPoolExecutor', 'ProcessPoolExecutor', ([], {'max_workers': 'WORKERS'}), '(max_workers=WORKERS)\n', (2871, 2892), False, 'from concurrent.futures import ProcessPoolExecutor, as_completed\n'), ((3631, 3707), 'codetiming.Timer', 'Timer', ([], {'name': '"""Bulk Index"""', 'text': '"""Validated data using Pydantic in {:.4f} sec"""'}), "(name='Bulk Index', text='Validated data using Pydantic in {:.4f} sec')\n", (3636, 3707), False, 'from codetiming import Timer\n'), ((3779, 3856), 'codetiming.Timer', 'Timer', ([], {'name': '"""Embed batches"""', 'text': '"""Created sentence embeddings in {:.4f} sec"""'}), "(name='Embed batches', text='Created sentence embeddings in {:.4f} sec')\n", (3784, 3856), False, 'from codetiming import Timer\n'), ((3981, 4050), 'codetiming.Timer', 'Timer', ([], {'name': '"""Create index"""', 'text': '"""Created IVF-PQ index in {:.4f} sec"""'}), "(name='Create index', text='Created IVF-PQ index in {:.4f} sec')\n", (3986, 4050), False, 'from codetiming import Timer\n'), ((3242, 3276), 'pandas.DataFrame.from_dict', 'pd.DataFrame.from_dict', (['embed_data'], {}), '(embed_data)\n', (3264, 3276), True, 'import pandas as pd\n'), ((3505, 3539), 'lancedb.pydantic.pydantic_to_schema', 'pydantic_to_schema', (['LanceModelWine'], {}), '(LanceModelWine)\n', (3523, 3539), False, 'from lancedb.pydantic import pydantic_to_schema\n'), ((1813, 1825), 'schemas.wine.Wine', 'Wine', ([], {}), '(**item)\n', (1817, 1825), False, 'from schemas.wine import LanceModelWine, Wine\n'), ((5060, 5074), 'pathlib.Path', 'Path', (['__file__'], {}), '(__file__)\n', (5064, 5074), False, 'from pathlib import Path\n'), ((462, 476), 'pathlib.Path', 'Path', (['__file__'], {}), '(__file__)\n', (466, 476), False, 'from pathlib import Path\n'), ((3185, 3206), 'concurrent.futures.as_completed', 'as_completed', (['futures'], {}), '(futures)\n', (3197, 3206), False, 'from concurrent.futures import ProcessPoolExecutor, as_completed\n')]