Update modeling_quasrav4.py

modeling_quasrav4.py

@@ -8,7 +8,7 @@ from transformers import PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.utils import logging
 
-from .configuration_quasrav4 import QuasraV4Config
+from .configuration_quasrav4 import InfinityFormerConfig
 
 logger = logging.get_logger(__name__)
 
@@ -41,7 +41,7 @@ class RotaryPositionEmbedding(nn.Module):
         return self.apply_rotary_pos_emb(x, cos.unsqueeze(0).unsqueeze(2), sin.unsqueeze(0).unsqueeze(2))
 
 class KernelFunction(nn.Module):
-    def __init__(self, config: QuasraV4Config):
+    def __init__(self, config: InfinityFormerConfig):
         super().__init__()
         self.kernel_type = config.kernel_type
         self.epsilon = config.kernel_epsilon
@@ -59,7 +59,7 @@ class KernelFunction(nn.Module):
             raise ValueError(f"Unknown kernel type: {self.kernel_type}")
 
 class GatedFeedForward(nn.Module):
-    def __init__(self, config: QuasraV4Config):
+    def __init__(self, config: InfinityFormerConfig):
         super().__init__()
         self.hidden_size = config.hidden_size
         self.intermediate_size = config.intermediate_size
@@ -80,7 +80,7 @@ class GatedFeedForward(nn.Module):
         return hidden_states + residual
 
 class LinearAttention(nn.Module):
-    def __init__(self, config: QuasraV4Config, layer_idx: int = 0):
+    def __init__(self, config: InfinityFormerConfig, layer_idx: int = 0):
         super().__init__()
         self.hidden_size = config.hidden_size
         self.num_heads = config.num_attention_heads
@@ -129,8 +129,8 @@ class LinearAttention(nn.Module):
 
 # --- Main Model Components ---
 
-class QuasraV4Layer(nn.Module):
-    def __init__(self, config: QuasraV4Config, layer_idx: int):
+class InfinityFormerLayer(nn.Module):
+    def __init__(self, config: InfinityFormerConfig, layer_idx: int):
         super().__init__()
         self.embed_dim = config.hidden_size
         self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
@@ -145,8 +145,8 @@ class QuasraV4Layer(nn.Module):
         hidden_states = self.ffn(hidden_states)
         return (hidden_states,)
 
-class QuasraV4Embeddings(nn.Module):
-    def __init__(self, config: QuasraV4Config):
+class InfinityFormerEmbeddings(nn.Module):
+    def __init__(self, config: InfinityFormerConfig):
         super().__init__()
         self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id if hasattr(config, 'pad_token_id') else 0)
         self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
@@ -166,11 +166,11 @@ class QuasraV4Embeddings(nn.Module):
         embeddings = self.dropout(embeddings)
         return embeddings
 
-class QuasraV4PreTrainedModel(PreTrainedModel):
-    config_class = QuasraV4Config
-    base_model_prefix = "model"
+class InfinityFormerPreTrainedModel(PreTrainedModel):
+    config_class = InfinityFormerConfig
+    base_model_prefix = "infinity_former"
     supports_gradient_checkpointing = True
-    _no_split_modules = ["QuasraV4Layer"]
+    _no_split_modules = ["InfinityFormerLayer"]
 
     def _init_weights(self, module):
         std = self.config.initializer_range
@@ -183,12 +183,12 @@ class QuasraV4PreTrainedModel(PreTrainedModel):
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
 
-class QuasraV4Model(QuasraV4PreTrainedModel):
-    def __init__(self, config: QuasraV4Config):
+class InfinityFormerModel(InfinityFormerPreTrainedModel):
+    def __init__(self, config: InfinityFormerConfig):
         super().__init__(config)
         self.config = config
-        self.embeddings = QuasraV4Embeddings(config)
-        self.layers = nn.ModuleList([QuasraV4Layer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+        self.embeddings = InfinityFormerEmbeddings(config)
+        self.layers = nn.ModuleList([InfinityFormerLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
         self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
         self.gradient_checkpointing = False
         self.post_init()
@@ -213,12 +213,12 @@ class QuasraV4Model(QuasraV4PreTrainedModel):
             return (hidden_states,)
         return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=None, hidden_states=None, attentions=None)
 
-class QuasraV4ForCausalLM(QuasraV4PreTrainedModel):
+class InfinityFormerForCausalLM(InfinityFormerPreTrainedModel):
     _auto_class = "AutoModelForCausalLM"
     
-    def __init__(self, config: QuasraV4Config):
+    def __init__(self, config: InfinityFormerConfig):
         super().__init__(config)
-        self.model = QuasraV4Model(config)
+        self.model = InfinityFormerModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.post_init()