Upload ConstBERT

modeling.py

@@ -6,6 +6,8 @@ from tqdm import tqdm
 from .colbert_configuration import ColBERTConfig
 from .tokenization_utils import QueryTokenizer, DocTokenizer
 import os
+
+
 class NullContextManager(object):
     def __init__(self, dummy_resource=None):
         self.dummy_resource = dummy_resource
@@ -54,7 +56,7 @@ class ConstBERT(BertPreTrainedModel):
     """
     _keys_to_ignore_on_load_unexpected = [r"cls"]
 
-    def __init__(self, config, colbert_config, verbose:int = 3):
+    def __init__(self, config, colbert_config, verbose:int = 0):
         super().__init__(config)
 
         self.config = config
@@ -175,7 +177,7 @@ class ConstBERT(BertPreTrainedModel):
 
                 return D
 
-    def queryFromText(self, queries, bsize=None, to_cpu=False, context=None, full_length_search=False):
+    def encode_query(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]
@@ -184,7 +186,7 @@ class ConstBERT(BertPreTrainedModel):
         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):
+    def encode_document(self, docs, bsize=None, keep_dims=True, to_cpu=False, showprogress=False, return_tokens=False):
         assert keep_dims in [True, False, 'flatten']
 
         if bsize: