code: import argparse
  import io
  import PIL
  import duckdb
  import lancedb
  import lance
  import pyarrow.compute as pc
  from transformers import CLIPModel, CLIPProcessor, CLIPTokenizerFast
  import gradio as gr
  
  MODEL_ID = None
  MODEL = None
  TOKENIZER = None
  PROCESSOR = None
  
  
  def create_table(dataset):
      db = lancedb.connect("~/datasets/demo")
      if "diffusiondb" in db.table_names():
          return db.open_table("diffusiondb")
  
      data = lance.dataset(dataset).to_table()
      tbl = db.create_table(
          "diffusiondb", data.filter(~pc.field("prompt").is_null()), mode="overwrite"
      )
      tbl.create_fts_index(["prompt"])
  
      return tbl
  
  
  def setup_clip_model(model_id):
      global MODEL_ID, MODEL, TOKENIZER, PROCESSOR
      MODEL_ID = model_id
      TOKENIZER = CLIPTokenizerFast.from_pretrained(MODEL_ID)
      MODEL = CLIPModel.from_pretrained(MODEL_ID)
      PROCESSOR = CLIPProcessor.from_pretrained(MODEL_ID)
  
  
  def embed_func(query):
      inputs = TOKENIZER([query], padding=True, return_tensors="pt")
      text_features = MODEL.get_text_features(**inputs)
      return text_features.detach().numpy()[0]
  
  
  def find_image_vectors(query):
      emb = embed_func(query)
      code = (
          "import lancedb\n"
          "db = lancedb.connect('~/datasets/demo')\n"
          "tbl = db.open_table('diffusiondb')\n\n"
          f"embedding = embed_func('{query}')\n"
          "tbl.search(embedding).limit(9).to_df()"
      )
      return (_extract(tbl.search(emb).limit(9).to_df()), code)
  
  
  def find_image_keywords(query):
      code = (
          "import lancedb\n"
          "db = lancedb.connect('~/datasets/demo')\n"
          "tbl = db.open_table('diffusiondb')\n\n"
          f"tbl.search('{query}').limit(9).to_df()"
      )
      return (_extract(tbl.search(query).limit(9).to_df()), code)
  
  
  def find_image_sql(query):
      code = (
          "import lancedb\n"
          "import duckdb\n"
          "db = lancedb.connect('~/datasets/demo')\n"
          "tbl = db.open_table('diffusiondb')\n\n"
          "diffusiondb = tbl.to_lance()\n"
          f"duckdb.sql('{query}').to_df()"
      )
      diffusiondb = tbl.to_lance()
      return (_extract(duckdb.sql(query).to_df()), code)
  
  
  def _extract(df):
      image_col = "image"
      return [
          (PIL.Image.open(io.BytesIO(row[image_col])), row["prompt"])
          for _, row in df.iterrows()
      ]
  
  
  def _extract(df):
      image_col = "image"
      return [
          (PIL.Image.open(io.BytesIO(row[image_col])), row["prompt"])
          for _, row in df.iterrows()
      ]
  
  
  def create_gradio_dash():
      with gr.Blocks() as demo:
          with gr.Row():
              with gr.Tab("Embeddings"):
                  vector_query = gr.Textbox(
                      value="portraits of a person", show_label=False
                  )
                  b1 = gr.Button("Submit")
              with gr.Tab("Keywords"):
                  keyword_query = gr.Textbox(value="ninja turtle", show_label=False)
                  b2 = gr.Button("Submit")
              with gr.Tab("SQL"):
                  sql_query = gr.Textbox(
                      value="SELECT * from diffusiondb WHERE image_nsfw >= 2 LIMIT 9",
                      show_label=False,
                  )
                  b3 = gr.Button("Submit")
          with gr.Row():
              code = gr.Code(label="Code", language="python")
          with gr.Row():
              gallery = gr.Gallery(
                  label="Found images", show_label=False, elem_id="gallery"
              ).style(columns=[3], rows=[3], object_fit="contain", height="auto")
  
          b1.click(find_image_vectors, inputs=vector_query, outputs=[gallery, code])
          b2.click(find_image_keywords, inputs=keyword_query, outputs=[gallery, code])
          b3.click(find_image_sql, inputs=sql_query, outputs=[gallery, code])
  
      demo.launch()
  
  
  def args_parse():
      parser = argparse.ArgumentParser()
      parser.add_argument("--model_id", type=str, default="openai/clip-vit-base-patch32")
      parser.add_argument("--dataset", type=str, default="rawdata.lance")
      return parser.parse_args()
  
  
  if __name__ == "__main__":
      args = args_parse()
      setup_clip_model(args.model_id)
      tbl = create_table(args.dataset)
      create_gradio_dash()
  
apis: lancedb.connect
extract_api: [((299, 333), 'lancedb.connect', 'lancedb.connect', (['"""~/datasets/demo"""'], {}), "('~/datasets/demo')\n", (314, 333), False, 'import lancedb\n'), ((759, 802), 'transformers.CLIPTokenizerFast.from_pretrained', 'CLIPTokenizerFast.from_pretrained', (['MODEL_ID'], {}), '(MODEL_ID)\n', (792, 802), False, 'from transformers import CLIPModel, CLIPProcessor, CLIPTokenizerFast\n'), ((815, 850), 'transformers.CLIPModel.from_pretrained', 'CLIPModel.from_pretrained', (['MODEL_ID'], {}), '(MODEL_ID)\n', (840, 850), False, 'from transformers import CLIPModel, CLIPProcessor, CLIPTokenizerFast\n'), ((867, 906), 'transformers.CLIPProcessor.from_pretrained', 'CLIPProcessor.from_pretrained', (['MODEL_ID'], {}), '(MODEL_ID)\n', (896, 906), False, 'from transformers import CLIPModel, CLIPProcessor, CLIPTokenizerFast\n'), ((3761, 3786), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (3784, 3786), False, 'import argparse\n'), ((2502, 2513), 'gradio.Blocks', 'gr.Blocks', ([], {}), '()\n', (2511, 2513), True, 'import gradio as gr\n'), ((432, 454), 'lance.dataset', 'lance.dataset', (['dataset'], {}), '(dataset)\n', (445, 454), False, 'import lance\n'), ((2536, 2544), 'gradio.Row', 'gr.Row', ([], {}), '()\n', (2542, 2544), True, 'import gradio as gr\n'), ((3183, 3191), 'gradio.Row', 'gr.Row', ([], {}), '()\n', (3189, 3191), True, 'import gradio as gr\n'), ((3212, 3252), 'gradio.Code', 'gr.Code', ([], {'label': '"""Code"""', 'language': '"""python"""'}), "(label='Code', language='python')\n", (3219, 3252), True, 'import gradio as gr\n'), ((3266, 3274), 'gradio.Row', 'gr.Row', ([], {}), '()\n', (3272, 3274), True, 'import gradio as gr\n'), ((2212, 2238), 'io.BytesIO', 'io.BytesIO', (['row[image_col]'], {}), '(row[image_col])\n', (2222, 2238), False, 'import io\n'), ((2379, 2405), 'io.BytesIO', 'io.BytesIO', (['row[image_col]'], {}), '(row[image_col])\n', (2389, 2405), False, 'import io\n'), ((2563, 2583), 'gradio.Tab', 'gr.Tab', (['"""Embeddings"""'], {}), "('Embeddings')\n", (2569, 2583), True, 'import gradio as gr\n'), ((2616, 2675), 'gradio.Textbox', 'gr.Textbox', ([], {'value': '"""portraits of a person"""', 'show_label': '(False)'}), "(value='portraits of a person', show_label=False)\n", (2626, 2675), True, 'import gradio as gr\n'), ((2735, 2754), 'gradio.Button', 'gr.Button', (['"""Submit"""'], {}), "('Submit')\n", (2744, 2754), True, 'import gradio as gr\n'), ((2772, 2790), 'gradio.Tab', 'gr.Tab', (['"""Keywords"""'], {}), "('Keywords')\n", (2778, 2790), True, 'import gradio as gr\n'), ((2824, 2874), 'gradio.Textbox', 'gr.Textbox', ([], {'value': '"""ninja turtle"""', 'show_label': '(False)'}), "(value='ninja turtle', show_label=False)\n", (2834, 2874), True, 'import gradio as gr\n'), ((2896, 2915), 'gradio.Button', 'gr.Button', (['"""Submit"""'], {}), "('Submit')\n", (2905, 2915), True, 'import gradio as gr\n'), ((2933, 2946), 'gradio.Tab', 'gr.Tab', (['"""SQL"""'], {}), "('SQL')\n", (2939, 2946), True, 'import gradio as gr\n'), ((2976, 3073), 'gradio.Textbox', 'gr.Textbox', ([], {'value': '"""SELECT * from diffusiondb WHERE image_nsfw >= 2 LIMIT 9"""', 'show_label': '(False)'}), "(value='SELECT * from diffusiondb WHERE image_nsfw >= 2 LIMIT 9',\n    show_label=False)\n", (2986, 3073), True, 'import gradio as gr\n'), ((3150, 3169), 'gradio.Button', 'gr.Button', (['"""Submit"""'], {}), "('Submit')\n", (3159, 3169), True, 'import gradio as gr\n'), ((2097, 2114), 'duckdb.sql', 'duckdb.sql', (['query'], {}), '(query)\n', (2107, 2114), False, 'import duckdb\n'), ((3298, 3367), 'gradio.Gallery', 'gr.Gallery', ([], {'label': '"""Found images"""', 'show_label': '(False)', 'elem_id': '"""gallery"""'}), "(label='Found images', show_label=False, elem_id='gallery')\n", (3308, 3367), True, 'import gradio as gr\n'), ((529, 547), 'pyarrow.compute.field', 'pc.field', (['"""prompt"""'], {}), "('prompt')\n", (537, 547), True, 'import pyarrow.compute as pc\n')]