Upload RecursiveGPT2Model.py with huggingface_hub

RecursiveGPT2Model.py

@@ -0,0 +1,27 @@
+from transformers import GPT2LMHeadModel, GPT2Config
+from shared_attention import convert_to_recursive
+import torch
+
+class RecursiveGPT2Config(GPT2Config):
+    model_type = "recursive_gpt2"
+    
+    def __init__(self, K=2, rank=8, **kwargs):
+        super().__init__(**kwargs)
+        self.K = K
+        self.rank = rank
+
+class RecursiveGPT2LMHeadModel(GPT2LMHeadModel):
+    config_class = RecursiveGPT2Config
+    
+    def __init__(self, config):
+        # Initialize as regular GPT2 first
+        super().__init__(config)
+        
+        # Apply recursive modifications
+        convert_to_recursive(self, K=config.K, rank=config.rank)
+        
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        # This ensures the recursive modifications are applied when loading
+        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        return model