Views
No views yet
1from transformers import T5ForConditionalGeneration,AutoTokenizer
2
3model = T5ForConditionalGeneration.from_pretrained('erfansadraiye/detoxify')
4tokenizer = AutoTokenizer.from_pretrained('erfansadraiye/detoxify')
5device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
6model.to(device)
7
8def test_model(text):
9 inputs = tokenizer.encode("Detoxify following sentence from bad words: " + text, return_tensors='pt', max_length=128, truncation=True, padding='max_length')
10 inputs = inputs.to(device)
11 outputs = model.generate(inputs, max_length=128, num_beams=4, temperature=0.7)
12 print("Output:", tokenizer.decode(outputs[0], skip_special_tokens=True))
13
14text = "Oh for fuck sake. This crazy moron is not even getting better."
15print("Original:", text)
16test_model(text)
17
18# output: Oh, this is not even getting better.
19
20