Spaces:
Paused
Paused
| import numpy as np | |
| import json | |
| from datetime import datetime | |
| from config import MEMORY_SIZE, LOCAL_MEMORY_SIZE | |
| class MemoryManager: | |
| def __init__(self): | |
| self.memory = np.memmap('jairvismkv_memory.bin', dtype='uint8', mode='w+', shape=(MEMORY_SIZE,)) | |
| self.local_memory_size = LOCAL_MEMORY_SIZE | |
| self.last_backup_date = None | |
| def store(self, input_data, output_data): | |
| data = json.dumps({ | |
| 'timestamp': datetime.now().isoformat(), | |
| 'input': input_data, | |
| 'output': output_data | |
| }) | |
| data_bytes = data.encode('utf-8') | |
| data_size = len(data_bytes) | |
| available_index = self.find_available_space(data_size) | |
| if available_index is not None: | |
| self.memory[available_index:available_index+data_size] = np.frombuffer(data_bytes, dtype='uint8') | |
| self.memory.flush() | |
| self.manage_local_memory() | |
| else: | |
| print("Not enough space in memory to store the data") | |
| def find_available_space(self, size): | |
| for i in range(0, len(self.memory) - size): | |
| if np.all(self.memory[i:i+size] == 0): | |
| return i | |
| return None | |
| def manage_local_memory(self): | |
| current_size = self.memory.size * self.memory.itemsize | |
| if current_size > self.local_memory_size: | |
| excess_size = current_size - self.local_memory_size | |
| self.memory = self.memory[:self.local_memory_size] | |
| self.memory.flush() | |
| print(f"Local memory optimized. Current size: {self.local_memory_size / (1024**3):.2f} GB") | |
| def retrieve(self, query): | |
| # Implement retrieval logic here | |
| # This could involve searching through the memory for relevant information | |
| pass | |
| def create_backup(self): | |
| current_date = datetime.now() | |
| if self.last_backup_date is None or (current_date - self.last_backup_date).days >= 30: | |
| backup_file = f"jairvismkv_backup_{current_date.strftime('%Y%m%d')}.bin" | |
| self.memory.flush() | |
| # Implement backup logic here (e.g., using Hugging Face's file storage) | |
| self.last_backup_date = current_date | |
| print(f"Backup created: {backup_file}") | |
| def load_from_backup(self, backup_file): | |
| # Implement logic to load memory from a backup file | |
| pass | |