Upload 4 files

artifact.metadata

@@ -0,0 +1,60 @@
+{
+    "query_token_id": "[unused0]",
+    "doc_token_id": "[unused1]",
+    "query_token": "[Q]",
+    "doc_token": "[D]",
+    "ncells": null,
+    "centroid_score_threshold": null,
+    "ndocs": null,
+    "load_index_with_mmap": false,
+    "index_path": null,
+    "index_bsize": 64,
+    "nbits": 1,
+    "kmeans_niters": 4,
+    "resume": false,
+    "similarity": "cosine",
+    "bsize": 2,
+    "accumsteps": 2,
+    "lr": 1e-5,
+    "maxsteps": 500000,
+    "save_every": null,
+    "warmup": 20000,
+    "warmup_bert": null,
+    "relu": false,
+    "nway": 64,
+    "use_ib_negatives": true,
+    "reranker": false,
+    "distillation_alpha": 1.0,
+    "ignore_scores": false,
+    "model_name": null,
+    "query_maxlen": 32,
+    "attend_to_mask_tokens": false,
+    "interaction": "colbert",
+    "dim": 128,
+    "doc_maxlen": 250,
+    "mask_punctuation": true,
+    "checkpoint": "bert-base-uncased",
+    "triples": ".\/examples.json",
+    "collection": ".\/collection.tsv",
+    "queries": ".\/queries.train.tsv",
+    "index_name": null,
+    "overwrite": false,
+    "root": "",
+    "experiment": "default",
+    "index_root": null,
+    "name": "2024-06\/30\/19.56.27",
+    "rank": 0,
+    "nranks": 8,
+    "amp": true,
+    "gpus": 8,
+    "avoid_fork_if_possible": false,
+    "meta": {
+        "hostname": "",
+        "git_branch": "main",
+        "git_hash": "",
+        "git_commit_datetime": "2024-06-26 12:50:22+00:00",
+        "current_datetime": "Jul 04, 2024 ;  3:51AM UTC (+0000)",
+        "cmd": "train.py",
+        "version": "colbert-v0.4"
+    }
+}

colbert_configuration.py

@@ -0,0 +1,412 @@
+from dataclasses import dataclass
+
+import __main__
+
+import os
+import ujson
+from huggingface_hub import hf_hub_download
+import dataclasses
+import datetime
+from typing import Any
+from dataclasses import dataclass, fields
+import socket
+import git
+import time
+import torch
+import sys
+
+def torch_load_dnn(path):
+    if path.startswith("http:") or path.startswith("https:"):
+        dnn = torch.hub.load_state_dict_from_url(path, map_location='cpu')
+    else:
+        dnn = torch.load(path, map_location='cpu')
+    
+    return dnn
+
+class dotdict(dict):
+    """
+    dot.notation access to dictionary attributes
+    Credit: derek73 @ https://stackoverflow.com/questions/2352181
+    """
+    __getattr__ = dict.__getitem__
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+def get_metadata_only():
+    args = dotdict()
+
+    args.hostname = socket.gethostname()
+    try:
+        args.git_branch = git.Repo(search_parent_directories=True).active_branch.name
+        args.git_hash = git.Repo(search_parent_directories=True).head.object.hexsha
+        args.git_commit_datetime = str(git.Repo(search_parent_directories=True).head.object.committed_datetime)
+    except git.exc.InvalidGitRepositoryError as e:
+        pass
+    args.current_datetime = time.strftime('%b %d, %Y ; %l:%M%p %Z (%z)')
+    args.cmd = ' '.join(sys.argv)
+
+    return args
+
+def timestamp(daydir=False):
+    format_str = f"%Y-%m{'/' if daydir else '-'}%d{'/' if daydir else '_'}%H.%M.%S"
+    result = datetime.datetime.now().strftime(format_str)
+    return result
+
+@dataclass
+class DefaultVal:
+    val: Any
+    
+    def __hash__(self):
+        return hash(repr(self.val))
+
+    def __eq__(self, other):
+        self.val == other.val
+
+@dataclass
+class RunSettings:
+    """
+        The defaults here have a special status in Run(), which initially calls assign_defaults(),
+        so these aren't soft defaults in that specific context.
+    """
+
+    overwrite: bool = DefaultVal(False)
+
+    root: str = DefaultVal(os.path.join(os.getcwd(), 'experiments'))
+    experiment: str = DefaultVal('default')
+
+    index_root: str = DefaultVal(None)
+    name: str = DefaultVal(timestamp(daydir=True))
+
+    rank: int = DefaultVal(0)
+    nranks: int = DefaultVal(1)
+    amp: bool = DefaultVal(True)
+
+    total_visible_gpus = torch.cuda.device_count()
+    gpus: int = DefaultVal(total_visible_gpus)
+
+    avoid_fork_if_possible: bool = DefaultVal(False)
+
+    @property
+    def gpus_(self):
+        value = self.gpus
+
+        if isinstance(value, int):
+            value = list(range(value))
+
+        if isinstance(value, str):
+            value = value.split(',')
+
+        value = list(map(int, value))
+        value = sorted(list(set(value)))
+
+        assert all(device_idx in range(0, self.total_visible_gpus) for device_idx in value), value
+
+        return value
+
+    @property
+    def index_root_(self):
+        return self.index_root or os.path.join(self.root, self.experiment, 'indexes/')
+
+    @property
+    def script_name_(self):
+        if '__file__' in dir(__main__):
+            cwd = os.path.abspath(os.getcwd())
+            script_path = os.path.abspath(__main__.__file__)
+            root_path = os.path.abspath(self.root)
+
+            if script_path.startswith(cwd):
+                script_path = script_path[len(cwd):]
+
+            else:
+                try:
+                    commonpath = os.path.commonpath([script_path, root_path])
+                    script_path = script_path[len(commonpath):]
+                except:
+                    pass
+
+
+            assert script_path.endswith('.py')
+            script_name = script_path.replace('/', '.').strip('.')[:-3]
+
+            assert len(script_name) > 0, (script_name, script_path, cwd)
+
+            return script_name
+
+        return 'none'
+
+    @property
+    def path_(self):
+        return os.path.join(self.root, self.experiment, self.script_name_, self.name)
+
+    @property
+    def device_(self):
+        return self.gpus_[self.rank % self.nranks]
+
+
+@dataclass
+class TokenizerSettings:
+    query_token_id: str = DefaultVal("[unused0]")
+    doc_token_id: str = DefaultVal("[unused1]")
+    query_token: str = DefaultVal("[Q]")
+    doc_token: str = DefaultVal("[D]")
+
+
+@dataclass
+class ResourceSettings:
+    checkpoint: str = DefaultVal(None)
+    triples: str = DefaultVal(None)
+    collection: str = DefaultVal(None)
+    queries: str = DefaultVal(None)
+    index_name: str = DefaultVal(None)
+
+
+@dataclass
+class DocSettings:
+    dim: int = DefaultVal(128)
+    doc_maxlen: int = DefaultVal(220)
+    mask_punctuation: bool = DefaultVal(True)
+
+
+@dataclass
+class QuerySettings:
+    query_maxlen: int = DefaultVal(32)
+    attend_to_mask_tokens : bool = DefaultVal(False)
+    interaction: str = DefaultVal('colbert')
+
+
+@dataclass
+class TrainingSettings:
+    similarity: str = DefaultVal('cosine')
+
+    bsize: int = DefaultVal(32)
+
+    accumsteps: int = DefaultVal(1)
+
+    lr: float = DefaultVal(3e-06)
+
+    maxsteps: int = DefaultVal(500_000)
+
+    save_every: int = DefaultVal(None)
+
+    resume: bool = DefaultVal(False)
+
+    ## NEW:
+    warmup: int = DefaultVal(None)
+
+    warmup_bert: int = DefaultVal(None)
+
+    relu: bool = DefaultVal(False)
+
+    nway: int = DefaultVal(2)
+
+    use_ib_negatives: bool = DefaultVal(False)
+
+    reranker: bool = DefaultVal(False)
+
+    distillation_alpha: float = DefaultVal(1.0)
+
+    ignore_scores: bool = DefaultVal(False)
+
+    model_name: str = DefaultVal(None) # DefaultVal('bert-base-uncased')
+
+@dataclass
+class IndexingSettings:
+    index_path: str = DefaultVal(None)
+
+    index_bsize: int = DefaultVal(64)
+
+    nbits: int = DefaultVal(1)
+
+    kmeans_niters: int = DefaultVal(4)
+
+    resume: bool = DefaultVal(False)
+
+    @property
+    def index_path_(self):
+        return self.index_path or os.path.join(self.index_root_, self.index_name)
+
+@dataclass
+class SearchSettings:
+    ncells: int = DefaultVal(None)
+    centroid_score_threshold: float = DefaultVal(None)
+    ndocs: int = DefaultVal(None)
+    load_index_with_mmap: bool = DefaultVal(False)
+
+
+@dataclass
+class CoreConfig:
+    def __post_init__(self):
+        """
+        Source: https://stackoverflow.com/a/58081120/1493011
+        """
+
+        self.assigned = {}
+
+        for field in fields(self):
+            field_val = getattr(self, field.name)
+
+            if isinstance(field_val, DefaultVal) or field_val is None:
+                setattr(self, field.name, field.default.val)
+
+            if not isinstance(field_val, DefaultVal):
+                self.assigned[field.name] = True
+    
+    def assign_defaults(self):
+        for field in fields(self):
+            setattr(self, field.name, field.default.val)
+            self.assigned[field.name] = True
+
+    def configure(self, ignore_unrecognized=True, **kw_args):
+        ignored = set()
+
+        for key, value in kw_args.items():
+            self.set(key, value, ignore_unrecognized) or ignored.update({key})
+
+        return ignored
+
+        """
+        # TODO: Take a config object, not kw_args.
+
+        for key in config.assigned:
+            value = getattr(config, key)
+        """
+
+    def set(self, key, value, ignore_unrecognized=False):
+        if hasattr(self, key):
+            setattr(self, key, value)
+            self.assigned[key] = True
+            return True
+
+        if not ignore_unrecognized:
+            raise Exception(f"Unrecognized key `{key}` for {type(self)}")
+
+    def help(self):
+        print(ujson.dumps(self.export(), indent=4))
+
+    def __export_value(self, v):
+        v = v.provenance() if hasattr(v, 'provenance') else v
+
+        if isinstance(v, list) and len(v) > 100:
+            v = (f"list with {len(v)} elements starting with...", v[:3])
+
+        if isinstance(v, dict) and len(v) > 100:
+            v = (f"dict with {len(v)} keys starting with...", list(v.keys())[:3])
+
+        return v
+
+    def export(self):
+        d = dataclasses.asdict(self)
+
+        for k, v in d.items():
+            d[k] = self.__export_value(v)
+
+        return d
+
+@dataclass
+class BaseConfig(CoreConfig):
+    @classmethod
+    def from_existing(cls, *sources):
+        kw_args = {}
+
+        for source in sources:
+            if source is None:
+                continue
+
+            local_kw_args = dataclasses.asdict(source)
+            local_kw_args = {k: local_kw_args[k] for k in source.assigned}
+            kw_args = {**kw_args, **local_kw_args}
+
+        obj = cls(**kw_args)
+
+        return obj
+
+    @classmethod
+    def from_deprecated_args(cls, args):
+        obj = cls()
+        ignored = obj.configure(ignore_unrecognized=True, **args)
+
+        return obj, ignored
+
+    @classmethod
+    def from_path(cls, name):
+        with open(name) as f:
+            args = ujson.load(f)
+
+            if "config" in args:
+                args = args["config"]
+
+        return cls.from_deprecated_args(
+            args
+        )  # the new, non-deprecated version functions the same at this level.
+
+    @classmethod
+    def load_from_checkpoint(cls, checkpoint_path):
+        if checkpoint_path.endswith(".dnn"):
+            dnn = torch_load_dnn(checkpoint_path)
+            config, _ = cls.from_deprecated_args(dnn.get("arguments", {}))
+
+            # TODO: FIXME: Decide if the line below will have any unintended consequences. We don't want to overwrite those!
+            config.set("checkpoint", checkpoint_path)
+
+            return config
+
+        try:
+            checkpoint_path = hf_hub_download(
+                repo_id=checkpoint_path, filename="artifact.metadata"
+            ).split("artifact")[0]
+        except Exception:
+            pass
+        loaded_config_path = os.path.join(checkpoint_path, "artifact.metadata")
+        if os.path.exists(loaded_config_path):
+            loaded_config, _ = cls.from_path(loaded_config_path)
+            loaded_config.set("checkpoint", checkpoint_path)
+
+            return loaded_config
+
+        return (
+            None  # can happen if checkpoint_path is something like 'bert-base-uncased'
+        )
+
+    @classmethod
+    def load_from_index(cls, index_path):
+        # FIXME: We should start here with initial_config = ColBERTConfig(config, Run().config).
+        # This should allow us to say initial_config.index_root. Then, below, set config = Config(..., initial_c)
+
+        # default_index_root = os.path.join(Run().root, Run().experiment, 'indexes/')
+        # index_path = os.path.join(default_index_root, index_path)
+
+        # CONSIDER: No more plan/metadata.json. Only metadata.json to avoid weird issues when loading an index.
+
+        try:
+            metadata_path = os.path.join(index_path, "metadata.json")
+            loaded_config, _ = cls.from_path(metadata_path)
+        except:
+            metadata_path = os.path.join(index_path, "plan.json")
+            loaded_config, _ = cls.from_path(metadata_path)
+
+        return loaded_config
+
+    def save(self, path, overwrite=False):
+        assert overwrite or not os.path.exists(path), path
+
+        with open(path, "w") as f:
+            args = self.export()  # dict(self.__config)
+            args["meta"] = get_metadata_only()
+            args["meta"]["version"] = "colbert-v0.4"
+            # TODO: Add git_status details.. It can't be too large! It should be a path that Runs() saves on exit, maybe!
+
+            f.write(ujson.dumps(args, indent=4) + "\n")
+
+    def save_for_checkpoint(self, checkpoint_path):
+        assert not checkpoint_path.endswith(
+            ".dnn"
+        ), f"{checkpoint_path}: We reserve *.dnn names for the deprecated checkpoint format."
+
+        output_config_path = os.path.join(checkpoint_path, "artifact.metadata")
+        self.save(output_config_path, overwrite=True)
+
+
+@dataclass
+class ColBERTConfig(RunSettings, ResourceSettings, DocSettings, QuerySettings, TrainingSettings,
+                    IndexingSettings, SearchSettings, BaseConfig, TokenizerSettings):
+    pass

tokenization_utils.py

@@ -0,0 +1,191 @@
+import torch
+from constbert.colbert_configuration import ColBERTConfig
+from transformers import AutoTokenizer
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def _split_into_batches(ids, mask, bsize):
+    batches = []
+    for offset in range(0, ids.size(0), bsize):
+        batches.append((ids[offset:offset+bsize], mask[offset:offset+bsize]))
+
+    return batches
+
+def _sort_by_length(ids, mask, bsize):
+    if ids.size(0) <= bsize:
+        return ids, mask, torch.arange(ids.size(0))
+
+    indices = mask.sum(-1).sort().indices
+    reverse_indices = indices.sort().indices
+
+    return ids[indices], mask[indices], reverse_indices
+
+class QueryTokenizer():
+    def __init__(self, config: ColBERTConfig, verbose: int = 3):
+        self.tok = AutoTokenizer.from_pretrained(config.checkpoint)
+        self.tok.base = config.checkpoint
+        self.verbose = verbose
+
+        self.config = config
+        self.query_maxlen = config.query_maxlen
+        self.background_maxlen = 512 - self.query_maxlen + 1  # FIXME: Make this configurable
+
+        self.Q_marker_token, self.Q_marker_token_id = config.query_token, self.tok.convert_tokens_to_ids(config.query_token_id)
+        self.cls_token, self.cls_token_id = self.tok.cls_token, self.tok.cls_token_id
+        self.sep_token, self.sep_token_id = self.tok.sep_token, self.tok.sep_token_id
+        self.mask_token, self.mask_token_id = self.tok.mask_token, self.tok.mask_token_id
+        self.pad_token,self.pad_token_id = self.tok.pad_token,self.tok.pad_token_id
+        self.used = False
+
+    def tokenize(self, batch_text, add_special_tokens=False):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        tokens = [self.tok.tokenize(x, add_special_tokens=False) for x in batch_text]
+
+        if not add_special_tokens:
+            return tokens
+
+        prefix, suffix = [self.cls_token, self.Q_marker_token], [self.sep_token]
+        tokens = [prefix + lst + suffix + [self.mask_token] * (self.query_maxlen - (len(lst)+3)) for lst in tokens]
+
+        return tokens
+
+    def encode(self, batch_text, add_special_tokens=False):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        ids = self.tok(batch_text, add_special_tokens=False).to(DEVICE)['input_ids']
+
+        if not add_special_tokens:
+            return ids
+
+        prefix, suffix = [self.cls_token_id, self.Q_marker_token_id], [self.sep_token_id]
+        ids = [prefix + lst + suffix + [self.mask_token_id] * (self.query_maxlen - (len(lst)+3)) for lst in ids]
+
+        return ids
+
+    def tensorize(self, batch_text, bsize=None, context=None, full_length_search=False):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        # add placehold for the [Q] marker
+        batch_text = ['. ' + x for x in batch_text]
+
+        # Full length search is only available for single inference (for now)
+        # Batched full length search requires far deeper changes to the code base
+        assert(full_length_search == False or (type(batch_text) == list and len(batch_text) == 1))
+
+        if full_length_search:
+            # Tokenize each string in the batch
+            un_truncated_ids = self.tok(batch_text, add_special_tokens=False).to(DEVICE)['input_ids']
+            # Get the longest length in the batch
+            max_length_in_batch = max(len(x) for x in un_truncated_ids)
+            # Set the max length
+            max_length = self.max_len(max_length_in_batch)
+        else:
+            # Max length is the default max length from the config
+            max_length = self.query_maxlen
+
+        obj = self.tok(batch_text, padding='max_length', truncation=True,
+                       return_tensors='pt', max_length=max_length).to(DEVICE)
+
+        ids, mask = obj['input_ids'], obj['attention_mask']
+
+        # postprocess for the [Q] marker and the [MASK] augmentation
+        ids[:, 1] = self.Q_marker_token_id
+        ids[ids == self.pad_token_id] = self.mask_token_id
+
+        if context is not None:
+            assert len(context) == len(batch_text), (len(context), len(batch_text))
+
+            obj_2 = self.tok(context, padding='longest', truncation=True,
+                            return_tensors='pt', max_length=self.background_maxlen).to(DEVICE)
+
+            ids_2, mask_2 = obj_2['input_ids'][:, 1:], obj_2['attention_mask'][:, 1:]  # Skip the first [SEP]
+
+            ids = torch.cat((ids, ids_2), dim=-1)
+            mask = torch.cat((mask, mask_2), dim=-1)
+
+        if self.config.attend_to_mask_tokens:
+            mask[ids == self.mask_token_id] = 1
+            assert mask.sum().item() == mask.size(0) * mask.size(1), mask
+
+        if bsize:
+            batches = _split_into_batches(ids, mask, bsize)
+            return batches
+        
+        if self.used is False:
+            self.used = True
+
+            firstbg = (context is None) or context[0]
+            if self.verbose > 1:
+                print()
+                print("#> QueryTokenizer.tensorize(batch_text[0], batch_background[0], bsize) ==")
+                print(f"#> Input: {batch_text[0]}, \t\t {firstbg}, \t\t {bsize}")
+                print(f"#> Output IDs: {ids[0].size()}, {ids[0]}")
+                print(f"#> Output Mask: {mask[0].size()}, {mask[0]}")
+                print()
+
+        return ids, mask
+
+    # Ensure that query_maxlen <= length <= 500 tokens
+    def max_len(self, length):
+        return min(500, max(self.query_maxlen, length))
+
+
+class DocTokenizer():
+    def __init__(self, config: ColBERTConfig):
+        self.tok = AutoTokenizer.from_pretrained(config.checkpoint)
+        self.tok.base = config.checkpoint
+
+        self.config = config
+        self.doc_maxlen = config.doc_maxlen
+
+        self.D_marker_token, self.D_marker_token_id = self.config.doc_token, self.tok.convert_tokens_to_ids(self.config.doc_token_id)
+        self.cls_token, self.cls_token_id = self.tok.cls_token, self.tok.cls_token_id
+        self.sep_token, self.sep_token_id = self.tok.sep_token, self.tok.sep_token_id
+
+    def tokenize(self, batch_text, add_special_tokens=False):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        tokens = [self.tok.tokenize(x, add_special_tokens=False).to(DEVICE) for x in batch_text]
+
+        if not add_special_tokens:
+            return tokens
+
+        prefix, suffix = [self.cls_token, self.D_marker_token], [self.sep_token]
+        tokens = [prefix + lst + suffix for lst in tokens]
+
+        return tokens
+
+    def encode(self, batch_text, add_special_tokens=False):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        ids = self.tok(batch_text, add_special_tokens=False).to(DEVICE)['input_ids']
+
+        if not add_special_tokens:
+            return ids
+
+        prefix, suffix = [self.cls_token_id, self.D_marker_token_id], [self.sep_token_id]
+        ids = [prefix + lst + suffix for lst in ids]
+
+        return ids
+
+    def tensorize(self, batch_text, bsize=None):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        # add placehold for the [D] marker
+        batch_text = ['. ' + x for x in batch_text]
+
+        obj = self.tok(batch_text, padding='max_length', truncation='longest_first',
+                       return_tensors='pt', max_length=self.doc_maxlen).to(DEVICE)
+
+        ids, mask = obj['input_ids'], obj['attention_mask']
+
+        # postprocess for the [D] marker
+        ids[:, 1] = self.D_marker_token_id
+
+        if bsize:
+            ids, mask, reverse_indices = _sort_by_length(ids, mask, bsize)
+            batches = _split_into_batches(ids, mask, bsize)
+            return batches, reverse_indices
+
+        return ids, mask