Views
No views yet
1from transformers import MBartForConditionalGeneration, AutoTokenizer
2
3model = MBartForConditionalGeneration.from_pretrained("s-nlp/mbart-detox-en-ru").cuda()
4tokenizer = AutoTokenizer.from_pretrained("facebook/mbart-large-50")
51def paraphrase(text, model, tokenizer, n=None, max_length="auto", beams=3):
2 texts = [text] if isinstance(text, str) else text
3 inputs = tokenizer(texts, return_tensors="pt", padding=True)["input_ids"].to(
4 model.device
5 )
6 if max_length == "auto":
7 max_length = inputs.shape[1] + 10
8
9 result = model.generate(
10 inputs,
11 num_return_sequences=n or 1,
12 do_sample=True,
13 temperature=1.0,
14 repetition_penalty=10.0,
15 max_length=max_length,
16 min_length=int(0.5 * max_length),
17 num_beams=beams,
18 forced_bos_token_id=tokenizer.lang_code_to_id[tokenizer.tgt_lang]
19 )
20 texts = [tokenizer.decode(r, skip_special_tokens=True) for r in result]
21
22 if not n and isinstance(text, str):
23 return texts[0]
24 return texts