File size: 935 Bytes
ca671b7 7537a46 ca671b7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
tokenizer = AutoTokenizer.from_pretrained("BEE-spoke-data/hf_slimpajama-6B-28672-BPE-forT5")
special_tokens_dict = {'additional_special_tokens': ['[R]', '[S]', '[X]', '[NTP]']}
tokenizer.add_special_tokens(special_tokens_dict)
model = AutoModelForSeq2SeqLM.from_pretrained("/workspace/nanoT5/logs/2024-10-20/18-25-17/amazingvince/ul3-base").to("cuda")
prompt = "[NTP] The "
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
# Add decoder_input_ids
# decoder_input_ids = torch.ones((inputs.input_ids.shape[0], 1), dtype=torch.long) * model.config.decoder_start_token_id
# Generate
generated_ids = model.generate(
**inputs,
# decoder_input_ids=decoder_input_ids,
max_new_tokens=20,
no_repeat_ngram_size=5
)
# Decode the output
generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
print(generated_text) |