README modified

README.md

@@ -40,21 +40,14 @@ tokenizer = T5Tokenizer.from_pretrained(raw_model)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 max_size = 10000
 
-def get_paraphrased_sentences(model, tokenizer, sentence, num_return_sequences=5, num_beams=5):
-  # tokenize the text to be form of a list of token IDs
-  inputs = tokenizer([sentence], truncation=True, padding="longest", return_tensors="pt").to(model.device)
-  # generate the paraphrased sentences
-  outputs = model.generate(
-    **inputs,
-    num_beams=num_beams,
-    num_return_sequences=num_return_sequences,
-    max_length=max_size
-  )
-  # decode the generated sentences using the tokenizer to get them back to text
-  return tokenizer.batch_decode(outputs, skip_special_tokens=True)
+def paraphrase(text, beams=100, grams=10, num_return_sequences=5):
+    x = tokenizer(text, return_tensors='pt', padding=True).to(model.device)
+    max_size = int(x.input_ids.shape[1] * 1.5 + 10)
+    out = model.generate(**x, encoder_no_repeat_ngram_size=grams, num_beams=beams, num_return_sequences=num_return_sequences, max_length=max_size)
+    return tokenizer.batch_decode(out, skip_special_tokens=True)
   
-# sentence = "Vorrei chiedervi la procedura per recuperare la chiave di accesso al mio profilo personale. L'ho persa e vorrei recuperarla."
-# get_paraphrased_sentences(model, tokenizer, sentence, num_beams=100, num_return_sequences=5)
+sentence = "Due amici si incontrano al bar per discutere del modo migliore di generare parafrasi."
+print(paraphrase(sentence))
 
 ```