| 76 | |
| 77 | |
| 78 | def tokenize_prompt(tokenizer: Tokenizer, text: str, context_length: int) -> List[int]: |
| 79 | sot_token = 49406 |
| 80 | eot_token = 49407 |
| 81 | |
| 82 | token_ids = list(tokenizer.encode(text).ids) |
| 83 | if not token_ids or token_ids[0] != sot_token: |
| 84 | token_ids = [sot_token] + token_ids |
| 85 | if token_ids[-1] != eot_token: |
| 86 | token_ids.append(eot_token) |
| 87 | |
| 88 | if len(token_ids) > context_length: |
| 89 | token_ids = token_ids[:context_length] |
| 90 | token_ids[-1] = eot_token |
| 91 | while len(token_ids) < context_length: |
| 92 | token_ids.append(0) |
| 93 | return token_ids |
| 94 | |
| 95 | |
| 96 | def main() -> None: |