feat: use local model

modelling_maincoder.py

@@ -81,13 +81,29 @@ class MaincoderRotaryEmbedding(nn.Module):
 
     def __init__(self, config: MaincoderConfig, device=None):
         super().__init__()
-        self.rope_type = "llama3" if config.rope_scaling is not None else "default"
+        self.rope_type = (
+            config.rope_scaling.get("rope_type", "default")
+            if isinstance(config.rope_scaling, dict)
+            else "default"
+        )
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        if self.rope_type == "default":
+            inv_freq, self.attention_scaling = self._compute_default_rope_parameters(device)
+        else:
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
 
+    def _compute_default_rope_parameters(self, device=None) -> tuple[torch.Tensor, float]:
+        inv_freq = 1.0 / (
+            self.config.rope_theta
+            ** (
+                torch.arange(0, self.config.head_dim, 2, dtype=torch.int64, device=device).float()
+                / self.config.head_dim
+            )
+        )
+        return inv_freq, 1.0
+
     @torch.no_grad()
     @dynamic_rope_update
     def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
@@ -188,7 +204,7 @@ class MaincoderAttention(nn.Module):
         position_embeddings: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         past_key_values: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         batch_size, seq_len, _ = hidden_states.shape
@@ -212,7 +228,7 @@ class MaincoderAttention(nn.Module):
 
         # Update KV cache
         if past_key_values is not None:
-            cache_kwargs = {"cache_position": cache_position}
+            cache_kwargs = {"cache_position": position_ids}
             key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
         # Attention
@@ -253,7 +269,7 @@ class MaincoderDecoderLayer(GradientCheckpointingLayer):
         attention_mask: Optional[torch.Tensor] = None,
         position_embeddings: Optional[torch.Tensor] = None,
         past_key_values: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> torch.Tensor:
         # Self Attention
@@ -264,7 +280,7 @@ class MaincoderDecoderLayer(GradientCheckpointingLayer):
             position_embeddings=position_embeddings,
             attention_mask=attention_mask,
             past_key_values=past_key_values,
-            cache_position=cache_position,
+            position_ids=position_ids,
             **kwargs,
         )
         hidden_states = residual + hidden_states
@@ -332,7 +348,6 @@ class MaincoderModel(MaincoderPreTrainedModel):
         past_key_values: Optional[Cache] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         use_cache: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, BaseModelOutputWithPast]:
         if (input_ids is None) ^ (inputs_embeds is not None):
@@ -344,24 +359,22 @@ class MaincoderModel(MaincoderPreTrainedModel):
         if use_cache and past_key_values is None:
             past_key_values = DynamicCache()
 
-        if cache_position is None:
+        if position_ids is None:
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
+            token_positions = torch.arange(
                 past_seen_tokens,
                 past_seen_tokens + inputs_embeds.shape[1],
                 device=inputs_embeds.device,
             )
-
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
+            position_ids = token_positions.unsqueeze(0)
 
         # Create causal mask
         causal_mask = create_causal_mask(
             config=self.config,
-            input_embeds=inputs_embeds,
+            inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
-            cache_position=cache_position,
             past_key_values=past_key_values,
+            position_ids=position_ids,
         )
 
         # Position embeddings
@@ -374,7 +387,7 @@ class MaincoderModel(MaincoderPreTrainedModel):
                 attention_mask=causal_mask,
                 position_embeddings=position_embeddings,
                 past_key_values=past_key_values,
-                cache_position=cache_position,
+                position_ids=position_ids,
                 **kwargs,
             )
 
@@ -389,7 +402,7 @@ class MaincoderModel(MaincoderPreTrainedModel):
 class MaincoderForCausalLM(MaincoderPreTrainedModel, GenerationMixin):
     """Maincoder model with a causal language modeling head."""
 
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
 
     def __init__(self, config: MaincoderConfig):
         super().__init__(config)
@@ -422,7 +435,6 @@ class MaincoderForCausalLM(MaincoderPreTrainedModel, GenerationMixin):
         inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
         use_cache: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
         logits_to_keep: Union[int, torch.Tensor] = 0,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, CausalLMOutputWithPast]:
@@ -454,7 +466,6 @@ class MaincoderForCausalLM(MaincoderPreTrainedModel, GenerationMixin):
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
             use_cache=use_cache,
-            cache_position=cache_position,
             **kwargs,
         )
 

run.py

@@ -1,26 +1,30 @@
+from pathlib import Path
+
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+MODEL_DIR = Path(__file__).resolve().parent
+
 model = AutoModelForCausalLM.from_pretrained(
-    "MengLinMaker/Maincoder-1B",
+    MODEL_DIR,
     torch_dtype="auto",
     device_map="auto",
     trust_remote_code=True,
 )
 tokenizer = AutoTokenizer.from_pretrained(
-    "MengLinMaker/Maincoder-1B",
+    MODEL_DIR,
     trust_remote_code=True,
 )
 
 # Code completion example
-prompt = '''def fibonacci(n: int) -> int:
-    """Return the n-th Fibonacci number."""
+prompt = '''"""Complete the fibonacci function in Python."""
+def fibonacci(n: int) -> int:
 '''
 
 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 outputs = model.generate(
     **inputs,
     max_new_tokens=256,
-    temperature=0.2,
+    temperature=0.5,
     do_sample=True,
 )
 print(tokenizer.decode(outputs[0], skip_special_tokens=True))