from transformers.models.llama.configuration_llama import LlamaConfig | |
class LlamaLandmarkConfig(LlamaConfig): | |
model_type = "llama_with_landmark" | |
def __init__( | |
self, | |
mem_id=32001, | |
mem_freq=50, | |
train_context_length=512, | |
**kwargs, | |
): | |
self.mem_id = mem_id | |
self.mem_freq = mem_freq | |
self.train_context_length = train_context_length | |
super().__init__(**kwargs) | |