| from langchain.llms.base import LLM |
| from langchain.memory import ConversationBufferWindowMemory |
| from transformers import GPT2TokenizerFast |
| from langchain.schema.messages import get_buffer_string |
|
|
| def get_num_tokens(text): |
| tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") |
| return len(tokenizer.tokenize(text)) |
|
|
| def get_memory_num_tokens(memory): |
| buffer = memory.chat_memory.messages |
| return sum([get_num_tokens(get_buffer_string([m])) for m in buffer]) |
|
|
| def validate_memory_len(memory, max_token_limit=2000): |
| buffer = memory.chat_memory.messages |
| curr_buffer_length = get_memory_num_tokens(memory) |
| if curr_buffer_length > max_token_limit: |
| while curr_buffer_length > max_token_limit: |
| buffer.pop(0) |
| curr_buffer_length = get_memory_num_tokens(memory) |
| return memory |
|
|
| if __name__ == '__main__': |
| tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") |
| text = '''Hi''' |
| print(len(tokenizer.tokenize(text))) |
