| 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) |