Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3
4device = 'cuda' if torch.cuda.is_available() else 'cpu'
5
6tokenizer = AutoTokenizer.from_pretrained("the-hir0/pegasus-detoxify")
7model = AutoModelForSeq2SeqLM.from_pretrained("the-hir0/pegasus-detoxify").to(device)
8
9input_ids = tokenizer(inference_request, return_tensors="pt", padding=True).input_ids.to(device)
10outputs = model.generate(input_ids=input_ids)
11tokenizer.batch_decode(outputs, skip_special_tokens=True,temperature=0)