code: import os
  import shutil
  import tempfile
  import lancedb
  import pandas as pd
  import numpy as np
  import duckdb
  from typing import Optional, List
  from .data_load import load_batches, load_df
  from .model import BaseEmbeddingModel
  from .db import LabelsDB
  from .settings import DEFAULT_TABLE_NAME, DB_BATCH_LOAD, DB_BATCH_SIZE
  
  
  duckdb.sql(
      """
      INSTALL sqlite;
      LOAD sqlite;
      """
  )
  
  
  class VectorDB:
      def __init__(
          self,
          db_path: str,
          db: lancedb.DBConnection,
          table: lancedb.table.Table,
          model: BaseEmbeddingModel,
          data_path: str,
      ) -> None:
          self.db = db
          self.model = model
          self.tbl = table
          self.data_path = data_path
          self.labelsdb_path = os.path.join(db_path, "labels.db")
          self.labelsdb = LabelsDB(self.labelsdb_path)
  
      @staticmethod
      def from_data_path(
          data_path: str,
          db_path: str,
          model: BaseEmbeddingModel,
          delete_existing=True,
          batch_load: bool = DB_BATCH_LOAD,
          batch_size: int = DB_BATCH_SIZE,
      ):
          db = lancedb.connect(db_path)
  
          table_name = DEFAULT_TABLE_NAME
          if delete_existing and table_name in db.table_names():
              print(f"Dropping existing database {table_name}...")
              db.drop_table(table_name)
              print("done.")
  
          if table_name in db.table_names():
              print(f'Opening existing table "{table_name}"...')
              tbl = db.open_table(table_name)
              return VectorDB(db_path, db, tbl, model, data_path)
          else:
              print(f'Creating table "{table_name}"...')
              if batch_load:
                  tbl = load_batches(db, table_name, data_path, model, batch_size)
              else:
                  tbl = load_df(db, table_name, data_path, model)
              return VectorDB(db_path, db, tbl, model, data_path)
  
      def count_rows(self) -> int:
          return len(self.tbl)
  
      def get(self, image_path: str) -> pd.DataFrame:
          lance_tbl = self.tbl.to_lance()
          return duckdb.sql(
              f"SELECT * FROM lance_tbl WHERE image_path='{image_path}';"
          ).to_df()
  
      def add_label(self, image_path: str, label: str):
          self.labelsdb.add(image_path=image_path, label=label)
  
      def remove_label(self, image_path: str, label: str):
          self.labelsdb.remove(image_path=image_path, label=label)
  
      def get_labels(self, image_path: Optional[str] = None) -> List[str]:
          return self.labelsdb.get(image_path=image_path)
  
      def get_label_counts(self) -> dict:
          return self.labelsdb.counts()
  
      def random_search(self, limit: int) -> pd.DataFrame:
          lance_tbl = self.tbl.to_lance()
          df_hits = duckdb.sql(
              f"""
              SELECT lance_tbl.*, grouped_labels.labels FROM lance_tbl
              LEFT OUTER JOIN (
                  SELECT image_path, list(label) AS labels FROM sqlite_scan('{self.labelsdb_path}', 'labels') GROUP BY image_path
              ) AS grouped_labels
              ON (lance_tbl.image_path = grouped_labels.image_path)
              USING SAMPLE {limit} ROWS;
          """
          ).to_df()
          df_hits["labels"] = df_hits["labels"].fillna("").apply(list)
          df_hits.drop(columns=["vector"], inplace=True)
          return df_hits
  
      def search_by_image_path(
          self, image_path: str, limit: int, exclude_labeled: bool
      ) -> pd.DataFrame:
          full_image_path = os.path.join(self.data_path, image_path)
          image_embedding = self.model.embed_image_path(full_image_path)
          df_hits = self.__vector_embedding_search(
              image_embedding,
              limit,
              exclude_image_path=image_path,
              exclude_labeled=exclude_labeled,
          )
          return df_hits
  
      def search_by_text(
          self, query_string: str, limit: int, exclude_labeled: bool
      ) -> pd.DataFrame:
          query_str_embedding = self.model.embed_text(query=query_string)
          df_hits = self.__vector_embedding_search(
              query_str_embedding, limit, exclude_labeled
          )
          return df_hits
  
      def __vector_embedding_search(
          self,
          embedding: np.ndarray,
          limit: int,
          exclude_labeled: bool,
          exclude_image_path: str = None,
      ) -> pd.DataFrame:
          if exclude_labeled:
              exclude_image_paths = set(self.labelsdb.get_image_paths())
          else:
              exclude_image_paths = set()
  
          if exclude_image_path is not None:
              exclude_image_paths.add(exclude_image_path)
  
          if len(exclude_image_paths) == 0:
              df_hits = self.tbl.search(embedding).limit(limit).to_df()
          else:
              exclude_image_paths_str = ",".join(
                  [f"'{image_path}'" for image_path in exclude_image_paths]
              )
              df_hits = (
                  self.tbl.search(embedding)
                  .where(f"image_path NOT IN ({exclude_image_paths_str})")
                  .limit(limit + len(exclude_image_paths))
                  .to_df()
              )
              df_hits = df_hits[0:limit]
  
          df_hits.drop(columns=["vector"], inplace=True)
          df_hits = self.__join_labels(left_table=df_hits)
          return df_hits
  
      def __join_labels(self, left_table: pd.DataFrame) -> pd.DataFrame:
          df_join = duckdb.sql(
              f"""
              SELECT left_table.*, grouped_labels.labels FROM left_table
              LEFT OUTER JOIN (
                  SELECT image_path, list(label) AS labels FROM sqlite_scan('{self.labelsdb_path}', 'labels') GROUP BY image_path
              ) AS grouped_labels
              ON (left_table.image_path = grouped_labels.image_path)
              ORDER BY left_table._distance ASC;
          """
          ).to_df()
          df_join["labels"] = df_join["labels"].fillna("").apply(list)
          return df_join
  
      def create_zip_labeled_binary_data(self, output_dir: str, filename: str) -> str:
          os.makedirs(output_dir, exist_ok=True)
          lance_tbl = self.tbl.to_lance()
          df = duckdb.sql(
              f"""
              SELECT lance_tbl.*, grouped_labels.labels FROM lance_tbl
              INNER JOIN (
                  SELECT image_path, label AS labels FROM sqlite_scan('{self.labelsdb_path}', 'labels') WHERE (label='relevant' OR label='irrelevant')
              ) AS grouped_labels
              ON (lance_tbl.image_path = grouped_labels.image_path);
          """
          ).to_df()
          df.drop(columns=["vector"], inplace=True)
  
          # Save df_hits to a CSV file in a temporary folder
          with tempfile.TemporaryDirectory() as tmpdir:
              csv_path = os.path.join(tmpdir, "data.csv")
              df.to_csv(csv_path, index=False)
  
              # Copy all image paths to the same temporary folder
              for image_path in df["image_path"]:
                  src_path = os.path.join(self.data_path, image_path)
                  dst_path = os.path.join(tmpdir, image_path)
                  print("Copying", src_path, "to", dst_path)
                  os.makedirs(os.path.dirname(dst_path), exist_ok=True)
                  shutil.copy(src_path, dst_path)
  
              # Create a zip file containing all files from the temporary folder
              if output_dir is None:
                  output_dir = tempfile.gettempdir()
              if filename.endswith(".zip"):
                  zip_path = os.path.join(output_dir, filename)
              else:
                  zip_path = os.path.join(output_dir, filename + ".zip")
              shutil.make_archive(zip_path[:-4], "zip", tmpdir)
  
              return zip_path
  
apis: lancedb.connect
extract_api: [((323, 383), 'duckdb.sql', 'duckdb.sql', (['"""\n    INSTALL sqlite;\n    LOAD sqlite;\n    """'], {}), '("""\n    INSTALL sqlite;\n    LOAD sqlite;\n    """)\n', (333, 383), False, 'import duckdb\n'), ((743, 777), 'os.path.join', 'os.path.join', (['db_path', '"""labels.db"""'], {}), "(db_path, 'labels.db')\n", (755, 777), False, 'import os\n'), ((1088, 1112), 'lancedb.connect', 'lancedb.connect', (['db_path'], {}), '(db_path)\n', (1103, 1112), False, 'import lancedb\n'), ((3438, 3478), 'os.path.join', 'os.path.join', (['self.data_path', 'image_path'], {}), '(self.data_path, image_path)\n', (3450, 3478), False, 'import os\n'), ((5927, 5965), 'os.makedirs', 'os.makedirs', (['output_dir'], {'exist_ok': '(True)'}), '(output_dir, exist_ok=True)\n', (5938, 5965), False, 'import os\n'), ((6543, 6572), 'tempfile.TemporaryDirectory', 'tempfile.TemporaryDirectory', ([], {}), '()\n', (6570, 6572), False, 'import tempfile\n'), ((6607, 6639), 'os.path.join', 'os.path.join', (['tmpdir', '"""data.csv"""'], {}), "(tmpdir, 'data.csv')\n", (6619, 6639), False, 'import os\n'), ((7474, 7523), 'shutil.make_archive', 'shutil.make_archive', (['zip_path[:-4]', '"""zip"""', 'tmpdir'], {}), "(zip_path[:-4], 'zip', tmpdir)\n", (7493, 7523), False, 'import shutil\n'), ((2056, 2127), 'duckdb.sql', 'duckdb.sql', (['f"""SELECT * FROM lance_tbl WHERE image_path=\'{image_path}\';"""'], {}), '(f"SELECT * FROM lance_tbl WHERE image_path=\'{image_path}\';")\n', (2066, 2127), False, 'import duckdb\n'), ((2723, 3125), 'duckdb.sql', 'duckdb.sql', (['f"""\n            SELECT lance_tbl.*, grouped_labels.labels FROM lance_tbl\n            LEFT OUTER JOIN (\n                SELECT image_path, list(label) AS labels FROM sqlite_scan(\'{self.labelsdb_path}\', \'labels\') GROUP BY image_path\n            ) AS grouped_labels\n            ON (lance_tbl.image_path = grouped_labels.image_path)\n            USING SAMPLE {limit} ROWS;\n        """'], {}), '(\n    f"""\n            SELECT lance_tbl.*, grouped_labels.labels FROM lance_tbl\n            LEFT OUTER JOIN (\n                SELECT image_path, list(label) AS labels FROM sqlite_scan(\'{self.labelsdb_path}\', \'labels\') GROUP BY image_path\n            ) AS grouped_labels\n            ON (lance_tbl.image_path = grouped_labels.image_path)\n            USING SAMPLE {limit} ROWS;\n        """\n    )\n', (2733, 3125), False, 'import duckdb\n'), ((5307, 5720), 'duckdb.sql', 'duckdb.sql', (['f"""\n            SELECT left_table.*, grouped_labels.labels FROM left_table\n            LEFT OUTER JOIN (\n                SELECT image_path, list(label) AS labels FROM sqlite_scan(\'{self.labelsdb_path}\', \'labels\') GROUP BY image_path\n            ) AS grouped_labels\n            ON (left_table.image_path = grouped_labels.image_path)\n            ORDER BY left_table._distance ASC;\n        """'], {}), '(\n    f"""\n            SELECT left_table.*, grouped_labels.labels FROM left_table\n            LEFT OUTER JOIN (\n                SELECT image_path, list(label) AS labels FROM sqlite_scan(\'{self.labelsdb_path}\', \'labels\') GROUP BY image_path\n            ) AS grouped_labels\n            ON (left_table.image_path = grouped_labels.image_path)\n            ORDER BY left_table._distance ASC;\n        """\n    )\n', (5317, 5720), False, 'import duckdb\n'), ((6019, 6399), 'duckdb.sql', 'duckdb.sql', (['f"""\n            SELECT lance_tbl.*, grouped_labels.labels FROM lance_tbl\n            INNER JOIN (\n                SELECT image_path, label AS labels FROM sqlite_scan(\'{self.labelsdb_path}\', \'labels\') WHERE (label=\'relevant\' OR label=\'irrelevant\')\n            ) AS grouped_labels\n            ON (lance_tbl.image_path = grouped_labels.image_path);\n        """'], {}), '(\n    f"""\n            SELECT lance_tbl.*, grouped_labels.labels FROM lance_tbl\n            INNER JOIN (\n                SELECT image_path, label AS labels FROM sqlite_scan(\'{self.labelsdb_path}\', \'labels\') WHERE (label=\'relevant\' OR label=\'irrelevant\')\n            ) AS grouped_labels\n            ON (lance_tbl.image_path = grouped_labels.image_path);\n        """\n    )\n', (6029, 6399), False, 'import duckdb\n'), ((6825, 6865), 'os.path.join', 'os.path.join', (['self.data_path', 'image_path'], {}), '(self.data_path, image_path)\n', (6837, 6865), False, 'import os\n'), ((6893, 6925), 'os.path.join', 'os.path.join', (['tmpdir', 'image_path'], {}), '(tmpdir, image_path)\n', (6905, 6925), False, 'import os\n'), ((7071, 7102), 'shutil.copy', 'shutil.copy', (['src_path', 'dst_path'], {}), '(src_path, dst_path)\n', (7082, 7102), False, 'import shutil\n'), ((7247, 7268), 'tempfile.gettempdir', 'tempfile.gettempdir', ([], {}), '()\n', (7266, 7268), False, 'import tempfile\n'), ((7338, 7372), 'os.path.join', 'os.path.join', (['output_dir', 'filename'], {}), '(output_dir, filename)\n', (7350, 7372), False, 'import os\n'), ((7418, 7461), 'os.path.join', 'os.path.join', (['output_dir', "(filename + '.zip')"], {}), "(output_dir, filename + '.zip')\n", (7430, 7461), False, 'import os\n'), ((7013, 7038), 'os.path.dirname', 'os.path.dirname', (['dst_path'], {}), '(dst_path)\n', (7028, 7038), False, 'import os\n')]