Upload memory_func.py

deal with memory capacity

models_for_langchain/memory_func.py

@@ -0,0 +1,26 @@
+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)))