code: from datasets import load_dataset
  from enum import Enum
  import lancedb
  from tqdm import tqdm
  from IPython.display import display
  import clip
  import torch
  
  
  class Animal(Enum):
      italian_greyhound = 0
      coyote = 1
      beagle = 2
      rottweiler = 3
      hyena = 4
      greater_swiss_mountain_dog = 5
      Triceratops = 6
      french_bulldog = 7
      red_wolf = 8
      egyption_cat = 9
      chihuahua = 10
      irish_terrier = 11
      tiger_cat = 12
      white_wolf = 13
      timber_wolf = 14
  
  
  def embed(img):
      image = preprocess(img).unsqueeze(0).to(device)
      embs = model.encode_image(image)
      return embs.detach().numpy()[0].tolist()
  
  
  def image_search(id):
      print("\n----- Image Search -----\n")
      print(Animal(test[id]["labels"]).name)
      display(test[id]["img"])
  
      res = tbl.search(embed(test[id]["img"])).limit(5).to_df()
      print(res)
      for i in range(5):
          print(Animal(res["label"][i]).name)
          data_id = int(res["id"][i])
          display(dataset[data_id]["img"])
  
  
  def embed_txt(txt):
      text = clip.tokenize([txt]).to(device)
      embs = model.encode_text(text)
      return embs.detach().numpy()[0].tolist()
  
  
  def text_search(text):
      print("\n----- Text Search -----\n")
      res = tbl.search(embed_txt(text)).limit(5).to_df()
      print(res)
      for i in range(len(res)):
          print(Animal(res["label"][i]).name)
          data_id = int(res["id"][i])
          display(dataset[data_id]["img"])
  
  
  def create_data(db):
      tbl = db.create_table(
          "animal_images",
          [{"vector": embed(dataset[0]["img"]), "id": 0, "label": dataset[0]["labels"]}],
      )
  
      data = []
      for i in tqdm(range(1, len(dataset))):
          data.append(
              {"vector": dataset[i]["img"], "id": i, "label": dataset[i]["labels"]}
          )
  
      batched_data = [data[n : n + 50] for n in range(0, len(data), 50)]
  
      for i in tqdm(batched_data):
          batch_data = []
          for j in i:
              row = {}
              row["vector"] = embed(j["vector"])
              row["id"] = j["id"]
              row["label"] = j["label"]
              batch_data.append(row)
          tbl.add(batch_data)
  
      return tbl
  
  
  if __name__ == "__main__":
      global dataset
      dataset = load_dataset(
          "CVdatasets/ImageNet15_animals_unbalanced_aug1", split="train"
      )
  
      device = "cuda" if torch.cuda.is_available() else "cpu"
      global model, preprocess
      model, preprocess = clip.load("ViT-B/32", device=device)
  
      db = lancedb.connect("./data/tables")
  
      # This function will take ~10 minutes, run if you don't have the data yet
      # tbl = create_data(db)
  
      # Run this to open the table for future runs
      tbl = db.open_table("animal_images")
      print(tbl.to_pandas())
  
      global test
      test = load_dataset(
          "CVdatasets/ImageNet15_animals_unbalanced_aug1", split="validation"
      )
  
      image_search(0)
      text_search("a full white dog")
  
apis: lancedb.connect
extract_api: [((748, 772), 'IPython.display.display', 'display', (["test[id]['img']"], {}), "(test[id]['img'])\n", (755, 772), False, 'from IPython.display import display\n'), ((1853, 1871), 'tqdm.tqdm', 'tqdm', (['batched_data'], {}), '(batched_data)\n', (1857, 1871), False, 'from tqdm import tqdm\n'), ((2196, 2272), 'datasets.load_dataset', 'load_dataset', (['"""CVdatasets/ImageNet15_animals_unbalanced_aug1"""'], {'split': '"""train"""'}), "('CVdatasets/ImageNet15_animals_unbalanced_aug1', split='train')\n", (2208, 2272), False, 'from datasets import load_dataset\n'), ((2401, 2437), 'clip.load', 'clip.load', (['"""ViT-B/32"""'], {'device': 'device'}), "('ViT-B/32', device=device)\n", (2410, 2437), False, 'import clip\n'), ((2448, 2480), 'lancedb.connect', 'lancedb.connect', (['"""./data/tables"""'], {}), "('./data/tables')\n", (2463, 2480), False, 'import lancedb\n'), ((2734, 2820), 'datasets.load_dataset', 'load_dataset', (['"""CVdatasets/ImageNet15_animals_unbalanced_aug1"""'], {'split': '"""validation"""'}), "('CVdatasets/ImageNet15_animals_unbalanced_aug1', split=\n    'validation')\n", (2746, 2820), False, 'from datasets import load_dataset\n'), ((962, 994), 'IPython.display.display', 'display', (["dataset[data_id]['img']"], {}), "(dataset[data_id]['img'])\n", (969, 994), False, 'from IPython.display import display\n'), ((1394, 1426), 'IPython.display.display', 'display', (["dataset[data_id]['img']"], {}), "(dataset[data_id]['img'])\n", (1401, 1426), False, 'from IPython.display import display\n'), ((2311, 2336), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (2334, 2336), False, 'import torch\n'), ((1028, 1048), 'clip.tokenize', 'clip.tokenize', (['[txt]'], {}), '([txt])\n', (1041, 1048), False, 'import clip\n')]