Update base_colbert.py

stanford-futuredata · Apr 20, 2023 · 43ae550 · 43ae550
1 parent 51c5230
commit 43ae550
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/colbert/modeling/base_colbert.py b/colbert/modeling/base_colbert.py
@@ -26,10 +26,12 @@ def __init__(self, name_or_path, colbert_config=None):
         try:
             HF_ColBERT = class_factory(self.name)
         except:
-            HF_ColBERT = class_factory('bert-base-uncased')
+            self.name = 'bert-base-uncased' # TODO: Double check that this is appropriate here in all cases
+            HF_ColBERT = class_factory(self.name)
 
-        assert self.name is not None
-        HF_ColBERT = class_factory(self.name)
+        # assert self.name is not None
+        # HF_ColBERT = class_factory(self.name)
+
         self.model = HF_ColBERT.from_pretrained(name_or_path, colbert_config=self.colbert_config)
         self.raw_tokenizer = AutoTokenizer.from_pretrained(name_or_path)