DeDeckerThomas commited on
Commit
cc84815
1 Parent(s): 1f1b128

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +2 -2
README.md CHANGED
@@ -48,7 +48,7 @@ Kulkarni, Mayank, Debanjan Mahata, Ravneet Arora, and Rajarshi Bhowmik. "Learnin
48
  # Model parameters
49
  from transformers import (
50
  Text2TextGenerationPipeline,
51
- BartForConditionalGeneration,
52
  AutoTokenizer,
53
  )
54
  import numpy as np
@@ -57,7 +57,7 @@ import numpy as np
57
  class KeyphraseGenerationPipeline(Text2TextGenerationPipeline):
58
  def __init__(self, model, keyphrase_sep_token=";", *args, **kwargs):
59
  super().__init__(
60
- model=BartForConditionalGeneration.from_pretrained(model),
61
  tokenizer=AutoTokenizer.from_pretrained(model),
62
  *args,
63
  **kwargs
 
48
  # Model parameters
49
  from transformers import (
50
  Text2TextGenerationPipeline,
51
+ AutoModelForSeq2SeqLM,
52
  AutoTokenizer,
53
  )
54
  import numpy as np
 
57
  class KeyphraseGenerationPipeline(Text2TextGenerationPipeline):
58
  def __init__(self, model, keyphrase_sep_token=";", *args, **kwargs):
59
  super().__init__(
60
+ model=AutoModelForSeq2SeqLM.from_pretrained(model),
61
  tokenizer=AutoTokenizer.from_pretrained(model),
62
  *args,
63
  **kwargs