Upload ConstBERT

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,29 @@
+{
+  "_name_or_path": "constbert/",
+  "architectures": [
+    "ConstBERT"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoModel": "modeling.ConstBERT"
+  },
+  "classifier_dropout": null,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.48.1",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e63f29e724efa1b9461cdc11af501c0c0fc09ac8c2c334ebf6e5dc4e45e422be
+size 438386000

modeling.py

@@ -0,0 +1,235 @@
+import torch.nn as nn
+from transformers import BertPreTrainedModel, BertModel, AutoTokenizer
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer
+from constbert.colbert_configuration import ColBERTConfig
+from constbert.tokenization_utils import QueryTokenizer, DocTokenizer
+
+
+class NullContextManager(object):
+    def __init__(self, dummy_resource=None):
+        self.dummy_resource = dummy_resource
+    def __enter__(self):
+        return self.dummy_resource
+    def __exit__(self, *args):
+        pass
+
+class MixedPrecisionManager():
+    def __init__(self, activated):
+        self.activated = activated
+
+        if self.activated:
+            self.scaler = torch.cuda.amp.GradScaler()
+
+    def context(self):
+        return torch.cuda.amp.autocast() if self.activated else NullContextManager()
+
+    def backward(self, loss):
+        if self.activated:
+            self.scaler.scale(loss).backward()
+        else:
+            loss.backward()
+
+    def step(self, colbert, optimizer, scheduler=None):
+        if self.activated:
+            self.scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(colbert.parameters(), 2.0, error_if_nonfinite=False)
+
+            self.scaler.step(optimizer)
+            self.scaler.update()
+        else:
+            torch.nn.utils.clip_grad_norm_(colbert.parameters(), 2.0)
+            optimizer.step()
+        
+        if scheduler is not None:
+            scheduler.step()
+
+        optimizer.zero_grad()
+
+class ConstBERT(BertPreTrainedModel):
+    """
+        Shallow wrapper around HuggingFace transformers. All new parameters should be defined at this level.
+
+        This makes sure `{from,save}_pretrained` and `init_weights` are applied to new parameters correctly.
+    """
+    _keys_to_ignore_on_load_unexpected = [r"cls"]
+
+    def __init__(self, config, colbert_config, verbose:int = 3):
+        super().__init__(config)
+
+        self.config = config
+        self.dim = colbert_config.dim
+        self.linear = nn.Linear(config.hidden_size, colbert_config.dim, bias=False)
+        self.doc_project = nn.Linear(colbert_config.doc_maxlen, 32, bias=False)
+        self.query_project = nn.Linear(colbert_config.query_maxlen, 64, bias=False)
+
+        self.query_tokenizer = QueryTokenizer(colbert_config, verbose=verbose)
+        self.doc_tokenizer = DocTokenizer(colbert_config)
+        self.amp_manager = MixedPrecisionManager(True)
+
+        self.raw_tokenizer = AutoTokenizer.from_pretrained(colbert_config.checkpoint)
+        self.pad_token = self.raw_tokenizer.pad_token_id
+        self.use_gpu = colbert_config.total_visible_gpus > 0
+
+        setattr(self,self.base_model_prefix, BertModel(config))
+
+        # if colbert_config.relu:
+        #     self.score_scaler = nn.Linear(1, 1)
+
+        self.init_weights()
+
+        # if colbert_config.relu:
+        #     self.score_scaler.weight.data.fill_(1.0)
+        #     self.score_scaler.bias.data.fill_(-8.0)
+
+    @property
+    def LM(self):
+        base_model_prefix = getattr(self, "base_model_prefix")
+        return getattr(self, base_model_prefix)
+
+
+    @classmethod
+    def from_pretrained(cls, name_or_path):
+        colbert_config = ColBERTConfig(name_or_path)
+        colbert_config = ColBERTConfig.from_existing(ColBERTConfig.load_from_checkpoint(name_or_path), colbert_config)
+        obj = super().from_pretrained(name_or_path, colbert_config=colbert_config)
+        obj.base = name_or_path
+
+        return obj
+
+    @staticmethod
+    def raw_tokenizer_from_pretrained(name_or_path):
+        obj = AutoTokenizer.from_pretrained(name_or_path)
+        obj.base = name_or_path
+
+        return obj
+    
+
+    def _query(self, input_ids, attention_mask):
+        input_ids, attention_mask = input_ids.to(self.device), attention_mask.to(self.device)
+        Q = self.bert(input_ids, attention_mask=attention_mask)[0]
+        # Q = Q.permute(0, 2, 1) #(64, 128,32)
+        # Q = self.query_project(Q) #(64, 128,8)
+        # Q = Q.permute(0, 2, 1) #(64,8,128)
+        Q = self.linear(Q)
+        # mask = torch.ones(Q.shape[0], Q.shape[1], device=self.device).unsqueeze(2).float()
+
+        mask = torch.tensor(self.mask(input_ids, skiplist=[]), device=self.device).unsqueeze(2).float()
+        Q = Q * mask
+
+        return torch.nn.functional.normalize(Q, p=2, dim=2)
+
+    def _doc(self, input_ids, attention_mask, keep_dims=True):
+        assert keep_dims in [True, False, 'return_mask']
+
+        input_ids, attention_mask = input_ids.to(self.device), attention_mask.to(self.device)
+        D = self.bert(input_ids, attention_mask=attention_mask)[0]
+        D = D.permute(0, 2, 1) #(64, 128,180)
+        D = self.doc_project(D) #(64, 128,16)
+        D = D.permute(0, 2, 1) #(64,16,128)
+        D = self.linear(D)
+        mask = torch.ones(D.shape[0], D.shape[1], device=self.device).unsqueeze(2).float()
+
+        # mask = torch.tensor(self.mask(input_ids, skiplist=self.skiplist), device=self.device).unsqueeze(2).float()
+        D = D * mask        
+        D = torch.nn.functional.normalize(D, p=2, dim=2)
+        if self.use_gpu:
+            D = D.half()
+
+        if keep_dims is False:
+            D, mask = D.cpu(), mask.bool().cpu().squeeze(-1)
+            D = [d[mask[idx]] for idx, d in enumerate(D)]
+
+        elif keep_dims == 'return_mask':
+            return D, mask.bool()
+
+        return D
+
+    def mask(self, input_ids, skiplist):
+        mask = [[(x not in skiplist) and (x != self.pad_token) for x in d] for d in input_ids.cpu().tolist()]
+        return mask
+
+    def query(self, *args, to_cpu=False, **kw_args):
+        with torch.no_grad():
+            with self.amp_manager.context():
+                Q = self._query(*args, **kw_args)
+                return Q.cpu() if to_cpu else Q
+
+    def doc(self, *args, to_cpu=False, **kw_args):
+        with torch.no_grad():
+            with self.amp_manager.context():
+                D = self._doc(*args, **kw_args)
+
+                if to_cpu:
+                    return (D[0].cpu(), *D[1:]) if isinstance(D, tuple) else D.cpu()
+
+                return D
+
+    def queryFromText(self, queries, bsize=None, to_cpu=False, context=None, full_length_search=False):
+        if bsize:
+            batches = self.query_tokenizer.tensorize(queries, context=context, bsize=bsize, full_length_search=full_length_search)
+            batches = [self.query(input_ids, attention_mask, to_cpu=to_cpu) for input_ids, attention_mask in batches]
+            return torch.cat(batches)
+
+        input_ids, attention_mask = self.query_tokenizer.tensorize(queries, context=context, full_length_search=full_length_search)
+        return self.query(input_ids, attention_mask)
+
+    def docFromText(self, docs, bsize=None, keep_dims=True, to_cpu=False, showprogress=False, return_tokens=False):
+        assert keep_dims in [True, False, 'flatten']
+
+        if bsize:
+            text_batches, reverse_indices = self.doc_tokenizer.tensorize(docs, bsize=bsize)
+
+            returned_text = []
+            if return_tokens:
+                returned_text = [text for batch in text_batches for text in batch[0]]
+                returned_text = [returned_text[idx] for idx in reverse_indices.tolist()]
+                returned_text = [returned_text]
+
+            keep_dims_ = 'return_mask' if keep_dims == 'flatten' else keep_dims
+            batches = [self.doc(input_ids, attention_mask, keep_dims=keep_dims_, to_cpu=to_cpu)
+                       for input_ids, attention_mask in tqdm(text_batches, disable=not showprogress)]
+
+            if keep_dims is True:
+                D = _stack_3D_tensors(batches)
+                return (D[reverse_indices], *returned_text)
+
+            elif keep_dims == 'flatten':
+                D, mask = [], []
+
+                for D_, mask_ in batches:
+                    D.append(D_)
+                    mask.append(mask_)
+
+                D, mask = torch.cat(D)[reverse_indices], torch.cat(mask)[reverse_indices]
+
+                doclens = mask.squeeze(-1).sum(-1).tolist()
+
+                D = D.view(-1, self.colbert_config.dim)
+                D = D[mask.bool().flatten()].cpu()
+
+                return (D, doclens, *returned_text)
+
+            assert keep_dims is False
+
+            D = [d for batch in batches for d in batch]
+            return ([D[idx] for idx in reverse_indices.tolist()], *returned_text)
+
+        input_ids, attention_mask = self.doc_tokenizer.tensorize(docs)
+        return self.doc(input_ids, attention_mask, keep_dims=keep_dims, to_cpu=to_cpu)
+
+def _stack_3D_tensors(groups):
+    bsize = sum([x.size(0) for x in groups])
+    maxlen = max([x.size(1) for x in groups])
+    hdim = groups[0].size(2)
+
+    output = torch.zeros(bsize, maxlen, hdim, device=groups[0].device, dtype=groups[0].dtype)
+
+    offset = 0
+    for x in groups:
+        endpos = offset + x.size(0)
+        output[offset:endpos, :x.size(1)] = x
+        offset = endpos
+
+    return output