Upload ConstBERT

modeling.py

@@ -178,6 +178,8 @@ class ConstBERT(BertPreTrainedModel):
                 return D
 
     def encode_query(self, queries, bsize=None, to_cpu=False, context=None, full_length_search=False):
+        if type(queries) == str:
+            queries = [queries]
         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]
@@ -187,6 +189,8 @@ class ConstBERT(BertPreTrainedModel):
         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):
+        if type(docs) == str:
+            docs = [docs]
         assert keep_dims in [True, False, 'flatten']
 
         if bsize: