from transformers import GPT2LMHeadModel

class SynapseGPT2LMHeadModel(GPT2LMHeadModel):
    def synapse_generate(self, tokenizer, question, return_trace=True, **kwargs):
        from synapse_controller import SynapseController
        controller = SynapseController(model=self, tokenizer=tokenizer)
        return controller.generate(question, return_trace=return_trace, **kwargs)