conditional_gen adjustments

src/Surveyor.py

@@ -146,6 +146,16 @@ class Surveyor:
             else:
                 self.ledtokenizer = AutoTokenizer.from_pretrained(ledmodel_name)
                 self.ledmodel = T5ForConditionalGeneration.from_pretrained(ledmodel_name).to(self.torch_device)
+
+            if 'led' in ledmodel_name:
+                self.ledtokenizer = LEDTokenizer.from_pretrained(ledmodel_name)
+                self.ledmodel = LEDForConditionalGeneration.from_pretrained(ledmodel_name).to(self.torch_device)
+            elif 't5' in ledmodel_name:
+                self.ledtokenizer = AutoTokenizer.from_pretrained(ledmodel_name)
+                self.ledmodel = T5ForConditionalGeneration.from_pretrained(ledmodel_name).to(self.torch_device)
+            elif 'bart' in ledmodel_name:
+                self.ledtokenizer = AutoTokenizer.from_pretrained(ledmodel_name)
+                self.ledmodel = BartForConditionalGeneration.from_pretrained(ledmodel_name).to(self.torch_device)
             self.ledmodel.eval()
             if not no_save_models:
                 self.ledmodel.save_pretrained(models_dir + "/ledmodel")
@@ -170,12 +180,15 @@ class Surveyor:
             self.summ_model.eval()
             self.model = Summarizer(custom_model=self.summ_model, custom_tokenizer=self.summ_tokenizer)
 
-            if 't5' not in ledmodel_name:
+            if 'led' in ledmodel_name:
                 self.ledtokenizer = LEDTokenizer.from_pretrained(ledmodel_name)
                 self.ledmodel = LEDForConditionalGeneration.from_pretrained(models_dir + "/ledmodel").to(self.torch_device)
-            else:
-                self.ledtokenizer = T5Tokenizer.from_pretrained(ledmodel_name)
+            elif 't5' in ledmodel_name:
+                self.ledtokenizer = AutoTokenizer.from_pretrained(ledmodel_name)
                 self.ledmodel = T5ForConditionalGeneration.from_pretrained(models_dir + "/ledmodel").to(self.torch_device)
+            elif 'bart' in ledmodel_name:
+                self.ledtokenizer = AutoTokenizer.from_pretrained(ledmodel_name)
+                self.ledmodel = BartForConditionalGeneration.from_pretrained(models_dir + "/ledmodel").to(self.torch_device)
             self.ledmodel.eval()
 
             self.embedder = SentenceTransformer(models_dir + "/embedder")