code clean up

.gitattributes

@@ -33,8 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-.bin filter=lfs diff=lfs merge=lfs -text
-.pt filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
 *.hdf5 filter=lfs diff=lfs merge=lfs -text
 *.index filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text

.vscode/settings.json

@@ -1,6 +0,0 @@
-{
-    "[python]": {
-        "editor.defaultFormatter": "ms-python.autopep8"
-    },
-    "python.formatting.provider": "none"
-}

app.py

@@ -1,6 +1,4 @@
 from pathlib import Path
-import datetime
-import json
 import os
 import time
 import gradio as gr
@@ -258,10 +256,11 @@ The service is a research preview intended for non-commercial use only, subject
 def build_demo():
     textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False)
     imagebox = gr.Image(type="pil")
-    state = gr.State()
 
     with gr.Blocks(title="GPT-K", theme=gr.themes.Base()) as demo:
+        state = gr.State()
         gr.Markdown(title_markdown)
+        
         with gr.Row():
             with gr.Column(scale=3):
                 gr.Examples(examples=[
@@ -274,10 +273,10 @@ def build_demo():
 
                 imagebox.render()
                 textbox.render()
-                with gr.Column():
+                with gr.Row():
                     submit_btn = gr.Button(value="📝 Submit")
-                    regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
-                    clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
+                    regenerate_btn = gr.Button(value="🔄 Regenerate", interactive=False)
+                    clear_btn = gr.Button(value="🗑️ Clear", interactive=False)
 
                 with gr.Accordion("Parameters", open=True):
                     with gr.Row():
@@ -289,7 +288,7 @@ def build_demo():
                     max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",)
 
             with gr.Column(scale=6):
-                chatbot = gr.Chatbot(elem_id="chatbot", label="LLaVA Chatbot", height=550)
+                chatbot = gr.Chatbot(elem_id="chatbot", label="GPT-K Chatbot", height=550)
             
                 gr.Markdown("Retrieved Knowledge")
                 knwl_img, knwl_txt = [], []
@@ -303,7 +302,7 @@ def build_demo():
                                     with gr.Column(scale=7):
                                         knwl_txt.append(gr.Markdown())
                 knwl_vis = knwl_img + knwl_txt
-
+        
         gr.Markdown(tos_markdown)
         gr.Markdown(learn_more_markdown)
 

conversation.py

@@ -1,6 +1,6 @@
 import dataclasses
 from enum import auto, Enum
-from typing import List, Tuple
+from typing import List
 
 
 class SeparatorStyle(Enum):
@@ -197,7 +197,6 @@ conv_gptk = Conversation(
     sep=""
 )
 
-
 conv_vicuna_v0 = Conversation(
     system="A chat between a curious human and an artificial intelligence assistant. "
            "The assistant gives helpful, detailed, and polite answers to the human's questions.",

knowledge/cluster.py

@@ -1,178 +0,0 @@
-import argparse
-from pathlib import Path
-import numpy as np
-from tqdm import tqdm
-import h5py
-import time
-
-import faiss
-import torch
-from pytorch_lightning import seed_everything
-
-import sys
-sys.path.append('.')
-from knowledge.text_db import TextDB
-from knowledge.utils import nn_search, build_faiss_index, refine_cosine
-
-
-UNSEEN = -2
-NOISE = -1
-
-
-def dbscan(X, faiss_index, device, eps=0.1, min_points=1, k=2048, bs=512):
-    neighbors = []
-    N = (len(X) - 1) // bs + 1
-    for i in tqdm(range(N), dynamic_ncols=True, desc="Find nearest neighbors", mininterval=1.0):
-        Xi = X[i*bs: (i+1)*bs]
-        _, I = faiss_index.search(Xi, k*2)
-        S, I = refine_cosine(X, Xi, I, device, k)
-
-        for sim, idx in zip(S, I):
-            dist = 1. - sim
-            neighbors.append(idx[dist < eps])
-        
-    cluster_id = 0
-    n_points = len(X)
-    labels = np.array([
-        NOISE if len(neighbors[i]) < min_points else UNSEEN
-        for i in range(n_points)
-    ])
-
-    with tqdm(total=n_points, dynamic_ncols=True, desc="DBSCAN clustering", mininterval=1.0) as pbar:
-        for i in range(n_points):
-            if labels[i] == UNSEEN:
-                seeds = np.array([i, ])
-                labels[seeds] = cluster_id
-
-                while len(seeds) > 0:
-                    neighbor_seeds = set()
-                    for s in seeds:
-                        n = neighbors[s]
-                        if len(n) > 0:
-                            l = np.array(list(set(labels[n])))
-                            l = l[np.logical_and(l >= 0, l != cluster_id)]
-                            for li in l:
-                                labels[labels == li] = cluster_id
-
-                            n = n[labels[n] == UNSEEN]
-                            neighbor_seeds.update(n)
-
-                    seeds = np.array(list(neighbor_seeds))
-                    if len(seeds) > 0:
-                        assert np.all(labels[seeds] == UNSEEN)
-                        labels[seeds] = cluster_id
-
-                cluster_id += 1
-            
-            pbar.set_postfix(num_clusters=cluster_id)
-            pbar.update()
-    
-    label_set = np.sort(list(set(labels)))
-    label_set = label_set[label_set >= 0]
-    labels_mapping = {l1: l2 for l2, l1 in enumerate(label_set)}
-    labels_mapping[-1] = -1
-    labels = np.array([labels_mapping[l] for l in labels])
-            
-    return labels
-
-
-def extract_clusters(feat, text, labels, faiss_index, device, k=128, bs=8192):
-    clusters = {}
-    for i, l in enumerate(tqdm(labels, dynamic_ncols=True, desc="Label each samples", mininterval=1.0)):
-        if l >= 0:
-            try:
-                clusters[l]["feat"] += feat[i].astype(np.float64)
-                clusters[l]["N"] += 1
-            except KeyError:
-                clusters[l] = {"feat": feat[i].astype(np.float64), "N": 1}
-    
-    cc = []
-    for l in tqdm(list(clusters.keys()), dynamic_ncols=True, desc="Compute cluster centers", mininterval=1.0):
-        c = clusters[l]["feat"]/clusters[l]["N"]
-        cc.append(c.astype(np.float32))
-    cc = np.stack(cc)
-    cc /= np.linalg.norm(cc, keepdims=True, axis=-1)
-
-    idx = []
-    N = (len(cc) - 1) // bs + 1
-    for i in tqdm(range(N), dynamic_ncols=True, desc="Find nearest neighbors", mininterval=1.0):
-        cc_i = cc[i*bs: (i+1)*bs]
-        _, I = faiss_index.search(cc_i, k)
-        _, I = refine_cosine(feat, cc_i, I, device, 1)
-        idx.append(I[:, 0])
-    idx = np.unique(np.concatenate(idx))
-    text = [text[i] for i in idx]
-    feat = np.stack([feat[i] for i in idx])
-
-    return feat, text
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Cluster knowledge database using DBSCAN")
-    parser.add_argument("--knowledge_db", type=str, required=True)
-    parser.add_argument("--seed", type=int, default=12345)
-    parser.add_argument("--eps", type=float, default=0.1)
-    parser.add_argument("--ms", type=int, default=1)
-    parser.add_argument("--ratio", type=float, default=None)
-    parser.add_argument("--device", type=int, default=None)
-    args = parser.parse_args()
-
-    # parse exp name
-    args.knowledge_db = Path(args.knowledge_db)
-    exp_name = args.knowledge_db.parent.name
-    exp_name += f"(dbscan)(eps-{args.eps})(ms-{args.ms})"
-    save_root = args.knowledge_db.parent.parent/exp_name
-    setattr(args, "save_root", save_root)
-    args.save_root.mkdir(parents=True, exist_ok=True)
-
-    args.device = torch.device("cuda", args.device) \
-        if args.device is not None else torch.device("cpu")
-
-    seed_everything(args.seed, workers=True)
-    print(args)
-    
-    # load feature, text, and faiss index from knowledge db
-    knowledge_db = TextDB(args.knowledge_db)
-    feat = knowledge_db.feature.astype(np.float32)
-    text = knowledge_db.text
-    if args.ratio is not None:
-        N = int(len(feat) * args.ratio)
-        feat, text = feat[:N], text[:N]
-    faiss_index = faiss.read_index(str(args.knowledge_db.parent/"faiss.index"))
-    print("Add data to faiss index...", end="\r")
-    ts = time.time()
-    faiss_index.add(feat)
-    print(f"Add data to faiss index...done in {time.time() - ts:.2f} secs")
-    
-    # DBSCAN clustering
-    labels_file = args.save_root/"labels.npy"
-    if labels_file.exists():
-        labels = np.load(labels_file)
-    else:
-        labels = dbscan(feat, faiss_index, args.device, args.eps, args.ms)
-        with open(labels_file, 'wb') as f:
-            np.save(f, labels)
-
-    # extract clusters
-    feat, text = extract_clusters(feat, text, labels, faiss_index, args.device)
-    with h5py.File(args.save_root/f"knowledge_db.hdf5", "w") as f:
-        bs = 65536
-        N = (len(feat) - 1) // bs + 1
-        for i in tqdm(range(N), dynamic_ncols=True, desc="Saving clustered DB", mininterval=1.0):
-            g = f.create_group(str(i))
-            g.create_dataset("feature", data=feat[i*bs: (i+1)*bs], compression="gzip")
-            g.create_dataset("text", data=text[i*bs: (i+1)*bs], compression="gzip")
-    
-    # build faiss index for the clustered DB
-    index = build_faiss_index(feat, gpus=[args.device.index, ])
-    faiss.write_index(index, str(args.save_root/"faiss.index"))
-
-    # some stats
-    noise_ratio = np.sum(labels == -1) / len(labels)
-    n_clusters, n_samples = len(text), len(labels)
-    msg = f"n_samples = {n_samples:,}; n_clusters = {n_clusters:,}; noise_ratio = {noise_ratio*100:.3f}%\n"
-    with open(save_root/"info.txt", "w") as f:
-        f.write(msg)
-    print(msg)
-
-

knowledge/retrieve.py

@@ -1,28 +1,8 @@
-import argparse
-from pathlib import Path
 import h5py
-import time
-import shutil
 import numpy as np
-import subprocess
-import time
 from tqdm import tqdm
-
-import faiss
-import open_clip
 import torch
-import torch.distributed as dist
-from torch.utils.data import DataLoader
-from pytorch_lightning import callbacks
-from pytorch_lightning import Trainer, LightningModule, seed_everything
-
-import sys
-sys.path.append('.')
-from dataset import coco, cc, llava
-from knowledge.utils import refine_cosine
-from knowledge import text_db
 from knowledge import TextDB
-from train.utils import ExpName
 
 
 class ImageCropsIdx:
@@ -123,205 +103,3 @@ class KnowAugImageCropsCombined:
             }
 
         return ret
-        
-
-class ImageCropsIdxBuilder(LightningModule):
-    def __init__(self, args, model: open_clip.model.CLIP):
-        super().__init__()
-
-        self.args = args
-        self.save_root = args.save_root
-        self.k = args.k
-        self.model = model
-        
-    def on_validation_epoch_start(self):
-        if self.global_rank == 0:
-            knowledge_db = TextDB(self.args.knowledge_db)
-            self.feature = knowledge_db.feature
-            self.text = knowledge_db.text
-
-            self.faiss_index = faiss.read_index(
-                str(Path(self.args.knowledge_db).parent/"faiss.index")
-            )
-            print("\nAdd data to faiss index...", end="\r")
-            ts = time.time()
-            self.faiss_index.add(self.feature)
-            print(f"Add data to faiss index...done in {time.time() - ts:.2f} secs")
-
-            with h5py.File(self.save_root/"knowledge_idx.hdf5", "a") as f:
-                f.attrs["fdim"] = self.feature.shape[-1]
-                f.attrs["file_hash"] = knowledge_db.file_hash
-        
-        self.trainer.strategy.barrier()
-    
-    def all_gather_object(self, data):
-        if self.trainer.world_size > 1:
-            gathered = [None for _ in range(self.trainer.world_size)]
-            dist.all_gather_object(gathered, data)
-            data = gathered
-        else:
-            data = [data, ]
-        
-        return data
-    
-    def broadcast_object(self, data, src_rank=0):
-        if self.trainer.world_size > 1:
-            if self.global_rank == src_rank:
-                data_list = [data, ] * self.trainer.world_size
-            else:
-                data_list = [None, ] * self.trainer.world_size
-            
-            dist.broadcast_object_list(data_list, src=src_rank)
-            return data_list[0]
-        else:
-            return data
-    
-    def search(self, images, topk):
-        query = self.model.encode_image(images, normalize=True)
-        query = query.cpu().numpy()
-        query = self.all_gather_object(query)
-        query = np.concatenate(query)
-
-        if self.global_rank == 0:
-            _, I = self.faiss_index.search(query, 4*topk)
-            S, I = refine_cosine(self.feature, query, I, self.device, topk)
-        else:
-            S = I = None
-
-        return S, I, query
-    
-    def validation_step(self, batch, batch_idx):
-        orig_imgs, five_imgs, nine_imgs, ids = batch
-
-        ids = ids.cpu().numpy()
-        ids = np.concatenate(self.all_gather_object(ids))
-
-        S_w, I_w, Q_w = self.search(orig_imgs, topk=self.k)
-        
-        S_f, I_f, Q_f = [], [], []
-        for i in range(five_imgs.shape[1]):
-            Si, Ii, Qi = self.search(five_imgs[:, i], topk=self.k)
-            S_f.append(Si)
-            I_f.append(Ii)
-            Q_f.append(Qi)
-        
-        S_n, I_n, Q_n = [], [], []
-        for i in range(nine_imgs.shape[1]):
-            Si, Ii, Qi = self.search(nine_imgs[:, i], topk=self.k)
-            S_n.append(Si)
-            I_n.append(Ii)
-            Q_n.append(Qi)
-
-        if self.global_rank == 0:
-            S_w, I_w, Q_w = np.expand_dims(S_w, axis=1), np.expand_dims(I_w, axis=1), np.expand_dims(Q_w, axis=1)
-            S_f, I_f, Q_f = np.stack(S_f, axis=1), np.stack(I_f, axis=1), np.stack(Q_f, axis=1)
-            S_n, I_n, Q_n = np.stack(S_n, axis=1), np.stack(I_n, axis=1), np.stack(Q_n, axis=1)
-
-            with h5py.File(self.save_root/"knowledge_idx.hdf5", "a") as f:
-                g = f.create_group(str(batch_idx))
-
-                g.create_dataset("image_ids", data=ids.astype(np.int32), compression="gzip")
-                
-                gw = g.create_group("whole")
-                gw.create_dataset("index", data=I_w.astype(np.int32), compression="gzip")
-                gw.create_dataset("score", data=S_w.astype(np.float32), compression="gzip")
-                gw.create_dataset("query", data=Q_w.astype(np.float32), compression="gzip")
-
-                gf = g.create_group("five")
-                gf.create_dataset("index", data=I_f.astype(np.int32), compression="gzip")
-                gf.create_dataset("score", data=S_f.astype(np.float32), compression="gzip")
-                gf.create_dataset("query", data=Q_f.astype(np.float32), compression="gzip")
-
-                gn = g.create_group("nine")
-                gn.create_dataset("index", data=I_n.astype(np.int32), compression="gzip")
-                gn.create_dataset("score", data=S_n.astype(np.float32), compression="gzip")
-                gn.create_dataset("query", data=Q_n.astype(np.float32), compression="gzip")
-
-    def on_validation_epoch_end(self):
-        if self.args.azcopy and self.global_rank == 0:
-            with open("azcopy/sas_output", "r") as f:
-                sas = f.readline()
-            sas_base, sas_key = sas.split("?")
-            sas = f"{sas_base}/knowledge_idx?{sas_key}"
-
-            cmd = ["azcopy/azcopy", "copy", str(self.args.save_root), sas, "--recursive=true"]
-            print(f"start copying data with command {cmd}")
-            ts = time.time()
-            subprocess.run(cmd)
-            print(f"done copying data in {time.time() - ts:.2f} secs")
-
-
-def main(args):
-    model, _, trans_img = open_clip.create_model_and_transforms(
-        args.clip_model, pretrained=text_db.CLIP_MODELS[args.clip_model]
-    )
-
-    print("load query dataset...")
-    if "coco" in args.query:
-        dset = coco.COCOImageCrops(Path(f"data/{args.query}"), trans=trans_img)
-        collate_crops = coco.collate_coco_crops
-    elif args.query == "cc3m":
-        dset = cc.CC3MImageCrops(Path("data/cc3m_instruct"), trans=trans_img)
-        collate_crops = cc.collate_cc_crops
-    elif args.query == "llava":
-        dset = llava.LLaVAImageCrops(Path("data/llava_bench"), trans=trans_img)
-        collate_crops = llava.collate_llava_crops
-    else:
-        raise ValueError
-    loader = DataLoader(
-        dset, batch_size=args.bs, shuffle=False, num_workers=args.num_workers,
-        drop_last=False, collate_fn=collate_crops
-    )
-
-    print("build model and trainer...")
-    pl_model = ImageCropsIdxBuilder(args, model)
-    model_summary = callbacks.RichModelSummary()
-    progress_bar = callbacks.TQDMProgressBar(args.refresh_rate)
-    trainer_callbacks = [model_summary, progress_bar]
-    trainer = Trainer(
-        sync_batchnorm=True,
-        precision=16,
-        accelerator='gpu',
-        devices=args.devices,
-        strategy="ddp",
-        default_root_dir=args.save_root,
-        callbacks=trainer_callbacks,
-        limit_val_batches=args.limit_val_batches
-    )
-
-    print("retrieve knowledge...")
-    trainer.validate(pl_model, dataloaders=loader)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description='Knowledge retrieval using image crops')
-    parser = Trainer.add_argparse_args(parser)
-    parser.add_argument('--query', type=str, choices=["coco14", "coco17", "cc3m", "llava"], required=True)
-    parser.add_argument('--knowledge_db', type=str, required=True)
-    parser.add_argument('--k', type=int, default=128)
-    parser.add_argument("--bs", type=int, default=128)
-    parser.add_argument("--num_workers", type=int, default=7)
-    parser.add_argument("--seed", type=int, default=12345)
-    parser.add_argument("--refresh_rate", type=int, default=1)
-    parser.add_argument("--azcopy", action="store_true")
-    args = parser.parse_args()
-
-    # parse exp_name
-    exp_name = ExpName(f"(query-{args.query})")
-    exp_name += Path(args.knowledge_db).parent.name
-    if args.azcopy:
-        setattr(args, "save_root", Path("azcopy")/str(exp_name))
-    else:
-        setattr(args, "save_root", Path("output")/"knowledge_idx"/str(exp_name))
-    shutil.rmtree(args.save_root, ignore_errors=True)
-    args.save_root.mkdir(parents=True, exist_ok=True)
-
-    # parse model
-    model = exp_name.get("clip-model")[1:-1]
-    model = model[len("clip-model-"):]
-    assert model in text_db.CLIP_MODELS.keys()
-    setattr(args, "clip_model", model)
-
-    print(args)
-    seed_everything(args.seed, workers=True)
-    main(args)

knowledge/text_db.py

@@ -1,26 +1,8 @@
-import argparse
-import itertools
-from pathlib import Path
-import shutil
 import h5py
-import time
-import subprocess
 from tqdm import tqdm
 import numpy as np
 import codecs
-
-import open_clip
-import faiss
-import torch
-import torch.distributed as dist
-from torch.utils.data import DataLoader
-from pytorch_lightning import callbacks
-from pytorch_lightning import Trainer, LightningModule, seed_everything
-
-import sys
-sys.path.append("./")
-from dataset import cc, words
-from knowledge.utils import file_hash, build_faiss_index
+from knowledge.utils import file_hash
 
 
 class TextDB:
@@ -59,139 +41,3 @@ class TextDB:
         
         return f, t
 
-
-class TextDBBuilder(LightningModule):
-    def __init__(self, args, model: open_clip.model.CLIP):
-        super().__init__()
-        self.args = args
-        self.model = model
-    
-    def validation_step(self, batch, batch_idx):
-        token, text = batch
-        feat = self.model.encode_text(token, normalize=True)
-
-        if self.trainer.world_size > 1:
-            text_gathered = [None for _ in range(self.trainer.world_size)]
-            dist.all_gather_object(text_gathered, text)
-            text = list(itertools.chain.from_iterable(text_gathered))
-
-            feat_gathered = [None for _ in range(self.trainer.world_size)]
-            dist.all_gather_object(feat_gathered, feat)
-            feat = torch.cat([x.to(self.device) for x in feat_gathered])
-            feat = feat.cpu().numpy()
-
-        if self.global_rank == 0:
-            with h5py.File(self.args.save_root/"knowledge_db.hdf5", "a") as f:
-                g = f.create_group(str(batch_idx))
-                g.create_dataset("feature", data=feat, compression="gzip")
-                g.create_dataset("text", data=text, compression="gzip")
-        
-    def validation_epoch_end(self, outputs):
-        if self.global_rank == 0:
-            knowledge_db = TextDB(self.args.save_root/"knowledge_db.hdf5")
-            feat = knowledge_db.feature
-
-            if self.args.devices == "-1":
-                num_devices = torch.cuda.device_count()
-                devices = list(range(num_devices))
-            else:
-                devices = [int(x) for x in args.devices.split(",") if x]
-            print(f"CUDA devices: {devices}")
-
-            index = build_faiss_index(feat, gpus=devices)
-            faiss.write_index(index, str(self.args.save_root/"faiss.index"))
-        self.trainer.strategy.barrier()
-
-        if self.args.azcopy and self.global_rank == 0:
-            with open("azcopy/sas_output", "r") as f:
-                sas = f.readline()
-            sas_base, sas_key = sas.split("?")
-            sas = f"{sas_base}/knowledge_db?{sas_key}"
-
-            cmd = ["azcopy/azcopy", "copy", str(self.args.save_root), sas, "--recursive=true"]
-            print(f"start copying data with command {cmd}")
-            ts = time.time()
-            subprocess.run(cmd)
-            print(f"done copying data in {time.time() - ts:.2f} secs")
-        self.trainer.strategy.barrier()
-
-
-DATASETS = {
-    "object": words.ObjsDataset,
-    "attribute": words.AttrsDataset,
-    "action": words.ActsDataset,
-    "cc3m": cc.CC3MTextDataset,
-    "cc12m": cc.CC12MTextDataset
-}
-
-
-def main(args):
-    model, _, _ = open_clip.create_model_and_transforms(
-        args.clip_model, pretrained=CLIP_MODELS[args.clip_model]
-    )
-    trans_txt = open_clip.get_tokenizer(args.clip_model)
-
-    print("load dataset...")
-    dset = DATASETS[args.dataset](Path(args.data_root), trans_txt)
-    loader = DataLoader(
-        dset, batch_size=args.bs, shuffle=False, num_workers=args.num_workers,
-        drop_last=False, collate_fn=cc.collate_cc_txt
-    )
-
-    print("build model and trainer...")
-    pl_model = TextDBBuilder(args, model)
-    model_summary = callbacks.RichModelSummary()
-    progress_bar = callbacks.TQDMProgressBar(args.refresh_rate)
-    trainer_callbacks = [model_summary, progress_bar]
-    trainer = Trainer(
-        sync_batchnorm=True,
-        precision=16,
-        accelerator='gpu',
-        devices=args.devices,
-        strategy="ddp",
-        default_root_dir=args.save_root,
-        callbacks=trainer_callbacks,
-        limit_val_batches=args.limit_val_batches
-    )
-
-    print("compute textual features...")
-    trainer.validate(pl_model, dataloaders=loader)
-
-
-CLIP_MODELS = {
-    'ViT-B-32': 'openai',
-    'ViT-B-16': 'openai',
-    'ViT-L-14': 'openai',
-    'ViT-g-14': 'laion2b_s34b_b88k',
-    'ViT-bigG-14': 'laion2b_s39b_b160k',
-    'convnext_xxlarge': 'laion2b_s34b_b82k_augreg_soup',
-}
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Build knowledge database of words")
-    parser = Trainer.add_argparse_args(parser)
-    parser.add_argument(
-        "--dataset", type=str, required=True, choices=["object", "attribute", "action", "cc3m", "cc12m"]
-    )
-    parser.add_argument("--data_root", type=str, default="data/conceptnet/conceptnet-assertions-5.7.0.csv")
-    parser.add_argument("--clip_model", type=str, default="ViT-g-14", choices=CLIP_MODELS.keys())
-    parser.add_argument("--bs", type=int, default=2**10)
-    parser.add_argument("--num_workers", type=int, default=7)
-    parser.add_argument("--seed", type=int, default=12345)
-    parser.add_argument("--refresh_rate", type=int, default=1)
-    parser.add_argument("--azcopy", action="store_true")
-    args = parser.parse_args()
-
-    # feature dir
-    exp_name = f"(dataset-{args.dataset})(clip-model-{args.clip_model})"
-    if args.azcopy:
-        setattr(args, "save_root", Path("azcopy")/"knowledge_db"/exp_name)
-    else:
-        setattr(args, "save_root", Path("output")/"knowledge_db"/exp_name)
-    shutil.rmtree(args.save_root, ignore_errors=True)
-    args.save_root.mkdir(parents=True, exist_ok=True)
-    
-    print(args)
-    seed_everything(args.seed, workers=True)
-    main(args)

knowledge/transforms.py

@@ -1,6 +1,5 @@
 import itertools
 from torchvision.transforms import functional as F
-import re
 
 
 def five_crop(image, ratio=0.6):
@@ -26,27 +25,3 @@ def nine_crop(image, ratio=0.4):
         images.append(F.crop(image, top, left, height, width))
 
     return images
-
-
-def pre_caption(caption, max_words=None):
-    # Ref: https://github.com/salesforce/LAVIS/blob/main/lavis/processors/blip_processors.py#L49-L68
-    caption = re.sub(
-        r"([.!\"()*#:;~])",
-        " ",
-        caption.lower(),
-    )
-    caption = re.sub(
-        r"\s{2,}",
-        " ",
-        caption,
-    )
-    caption = caption.rstrip("\n")
-    caption = caption.strip(" ")
-
-    # truncate caption
-    caption_words = caption.split(" ")
-    if max_words is not None and len(caption_words) > max_words:
-        caption = " ".join(caption_words[: max_words])
-
-    return caption
-    

knowledge/utils.py

@@ -1,11 +1,5 @@
-from tqdm import tqdm
 import numpy as np
-import time
-import math
-import bisect
 import hashlib
-import faiss
-from faiss import StandardGpuResources, index_cpu_to_gpu_multiple_py
 import torch
 
 
@@ -21,78 +15,6 @@ def file_hash(file):
     return hash_fn.hexdigest()
 
 
-def build_faiss_index(x, gpus=None):
-    # Ref: https://github.com/facebookresearch/faiss/wiki/Guidelines-to-choose-an-index
-    # Ref: https://gist.github.com/mdouze/46d6bbbaabca0b9778fca37ed2bcccf6
-
-    N, dim = x.shape
-    secs = [2**i for i in range(1, 15)]
-    d = secs[bisect.bisect_right(secs, dim) - 1] // 2
-    m = d // 4
-
-    if N <= 60000:
-        index_factory = "Flat"
-    elif N <= 2555904:
-        index_factory = f"IVF{int(8*math.sqrt(N))},Flat"
-    elif N <= 10223616:
-        index_factory = f"OPQ{m}_{d},IVF65536_HNSW32,PQ{m}x4fsr"
-    elif N <= 1e8:
-        index_factory = f"OPQ{m}_{d},IVF262144_HNSW32,PQ{m}x4fsr"
-    else:
-        index_factory = f"OPQ{m}_{d},IVF1048576_HNSW32,PQ{m}x4fsr"
-    print(f"train {index_factory} index on {N:,} x {dim} data")
-
-    index = faiss.index_factory(dim, index_factory)
-    if gpus is not None and N > 60000:
-        index_ivf = faiss.extract_index_ivf(index)
-        res = []
-        for _ in gpus:
-            r = StandardGpuResources()
-            r.noTempMemory()
-            res.append(r)
-        clustering_index = index_cpu_to_gpu_multiple_py(
-            res, faiss.IndexFlatL2(index_ivf.d), None, gpus
-        )
-        index_ivf.clustering_index = clustering_index
-
-    print("train index...", end="\r")
-    ts = time.time()
-    # commented out for index_factory = "Flat"
-    # assert not index.is_trained
-    index.train(x)
-    assert index.is_trained
-    print(f"train index...done in {time.time() - ts:.2f} secs")
-
-    index.nprobe = 64
-    index.quantizer_efSearch = 32
-
-    return index
-
-
-def nn_search(query, index, topk, bs=256, desc=None, disable_tqdm=True):
-    idx, dist = [], []
-    N = (len(query) - 1) // bs + 1
-    for i in tqdm(range(N), dynamic_ncols=True, desc=desc, disable=disable_tqdm):
-        D, I = index.search(query[i*bs: (i+1)*bs], topk)
-        idx.append(I)
-        dist.append(D)
-    idx = np.concatenate(idx)
-    dist = np.concatenate(dist)
-
-    return idx, dist
-
-
-def radius_search(query, index, r, bs=256, desc=None, disable_tqdm=True):
-    idx, dist = [], []
-    N = (len(query) - 1) // bs + 1
-    for i in tqdm(range(N), dynamic_ncols=True, desc=desc, disable=disable_tqdm):
-        L, D, I = index.range_search(query[i*bs: (i+1)*bs], r)
-        idx.extend([I[L[j]:L[j+1]] for j in range(len(L)-1)])
-        dist.extend([D[L[j]:L[j+1]] for j in range(len(L)-1)])
-
-    return idx, dist
-
-
 @torch.no_grad()
 def refine_cosine(Xa, Xq, I, device, k=None):
     if k is not None:
@@ -114,14 +36,3 @@ def refine_cosine(Xa, Xq, I, device, k=None):
     S_refined = np.stack(S_refined)
     
     return S_refined, I_refined
-
-
-def test_nn_search():
-    key = np.random.random((3000000, 512)).astype(np.float32)
-    key /= np.linalg.norm(key, keepdims=True, axis=1)
-    index = build_faiss_index(key, -1)
-
-    query = np.random.random((100000, 512)).astype(np.float32)
-    query /= np.linalg.norm(query, keepdims=True, axis=1)
-    idx_r = nn_search(query, index, r=0.5)
-    idx_k = nn_search(query, index, topk=10)

model/.gitattributes

@@ -1,2 +0,0 @@
-*.hdf5 filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text