Update src/paraphrase/Paraphrase.py

src/paraphrase/Paraphrase.py

@@ -1,15 +1,11 @@
-from nltk.tokenize import sent_tokenize
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
 import torch
 import src.exception.Exception.Exception as ExceptionCustom
 
 METHOD = "PARAPHRASE"
 
-tokenizer = AutoTokenizer.from_pretrained("t5-base")
-model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
-
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# model.to(device)
+tokenizer = PegasusTokenizer.from_pretrained('google/pegasus-xsum')
+model = PegasusForConditionalGeneration.from_pretrained('google/pegasus-xsum')
 
 def paraphraseParaphraseMethod(requestValue : str):
     exception = ""
@@ -24,20 +20,16 @@ def paraphraseParaphraseMethod(requestValue : str):
     for SENTENCE in tokenized_sent_list:
         text = "paraphrase: " + SENTENCE
 
-        encoding = tokenizer.encode_plus(text, pad_to_max_length=True, return_tensors="pt")
+        encoding = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
         input_ids, attention_masks = encoding["input_ids"], encoding["attention_mask"]
 
         beam_outputs = model.generate(
             input_ids=input_ids, 
             attention_mask=attention_masks,
-            do_sample=True,
             max_length=512,
-            top_k=120,
-            top_p=0.90,
-            early_stopping=False,
-            num_return_sequences=1,
-            no_repeat_ngram_size=2,
-            num_beams=1
+            num_beams=5,
+            length_penalty=0.8,
+            early_stopping=True
         )
 
         for beam_output in beam_outputs: