File size: 383 Bytes
b399ea6
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9

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)