Upload ConstBERT

modeling.py

@@ -177,7 +177,7 @@ class ConstBERT(BertPreTrainedModel):
 
                 return D
 
-    def encode_query(self, queries, bsize=None, to_cpu=False, context=None, full_length_search=False):
+    def encode_queries(self, queries, bsize=None, to_cpu=False, context=None, full_length_search=False):
         if type(queries) == str:
             queries = [queries]
         if bsize:
@@ -188,7 +188,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 encode_document(self, docs, bsize=None, keep_dims=True, to_cpu=False, showprogress=False, return_tokens=False):
+    def encode_documents(self, docs, bsize=None, keep_dims=True, to_cpu=False, showprogress=False, return_tokens=False):
         if type(docs) == str:
             docs = [docs]
         assert keep_dims in [True, False, 'flatten']