custom-decoder-ats / README.md
josh-oo's picture
Create README.md
96fb662
|
raw
history blame
1.41 kB
import torch
from transformers import AutoTokenizer
from transformers import AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("josh-oo/custom-decoder-ats")

##gerpt

#model = AutoModelForSeq2SeqLM.from_pretrained("josh-oo/custom-decoder-ats", trust_remote_code=True, revision="35197269f0235992fcc6b8363ca4f48558b624ff")
#decoder_tokenizer = AutoTokenizer.from_pretrained("josh-oo/gerpt2")

##dbmdz

model = AutoModelForSeq2SeqLM.from_pretrained("josh-oo/custom-decoder-ats", trust_remote_code=True, revision="4accedbe0b57d342d95ff546b6bbd3321451d504")
decoder_tokenizer = AutoTokenizer.from_pretrained("josh-oo/german-gpt2-easy")
decoder_tokenizer.add_tokens(['<</s>>','<<s>>','<<pad>>'])

##

example_text = "In tausenden Schweizer Privathaushalten kümmern sich Haushaltsangestellte um die Wäsche, betreuen die Kinder und sorgen für Sauberkeit. Durchschnittlich bekommen sie für die Arbeit rund 30 Franken pro Stunde Bruttolohn. Der grösste Teil von ihnen erhält aber 28 Franken."

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

test_input = tokenizer([example_text], return_tensors="pt", padding=True, pad_to_multiple_of=1024)
for key, value in test_input.items():
  test_input[key] = value.to(device)

outputs = model.generate(**test_input, num_beams=3, max_length=1024)
decoder_tokenizer.batch_decode(outputs)