Update README.md

README.md

@@ -48,7 +48,7 @@ Kulkarni, Mayank, Debanjan Mahata, Ravneet Arora, and Rajarshi Bhowmik. "Learnin
 # Model parameters
 from transformers import (
     Text2TextGenerationPipeline,
-    BartForConditionalGeneration,
+    AutoModelForSeq2SeqLM,
     AutoTokenizer,
 )
 import numpy as np
@@ -57,7 +57,7 @@ import numpy as np
 class KeyphraseGenerationPipeline(Text2TextGenerationPipeline):
     def __init__(self, model, keyphrase_sep_token=";", *args, **kwargs):
         super().__init__(
-            model=BartForConditionalGeneration.from_pretrained(model),
+            model=AutoModelForSeq2SeqLM.from_pretrained(model),
             tokenizer=AutoTokenizer.from_pretrained(model),
             *args,
             **kwargs