Upload folder using huggingface_hub

README.md

@@ -1,3 +1,122 @@
----
-license: mit
----
+---
+library_name: transformers
+tags:
+- pytorch
+- custom-implementation
+- graph-prediction
+- edge-prediction
+---
+
+# Llama Edge Prediction Model
+
+This repository contains a custom Llama 3 based model for edge prediction tasks. It predicts edge targets based on context IDs.
+
+## Model Description
+
+The model corresponds to a `Llama3` architecture with the following configuration:
+- Standard Llama 3 8B params (dim 4096, 32 layers, 32 heads, 8 KV heads)
+- Adjusted vocab size: 9942 (custom embeddings)
+- Intermediate size: 14336
+
+It uses a `UnifiedIdMapper` to map between original IDs (nodes/edges) and internal model IDs.
+
+## Repository Structure
+
+- `configuration_llama_edge.py`: Defines `LlamaEdgeConfig` (inherits from `PretrainedConfig`).
+- `modeling_llama_edge.py`: Defines `LlamaEdgeForCausalLM` and components (inherits from `PreTrainedModel`).
+- `id_mapper.py`: `UnifiedIdMapper` for ID mapping logic.
+- `inference.py`: Example script to run inference using the model and mapper.
+- `model.safetensors`: Model weights (required).
+- `unified_id_mapper.json`: Mapping data (required).
+
+## Usage
+
+### Loading the Model
+
+You can load the model using the provided classes:
+
+```python
+import torch
+from configuration_llama_edge import LlamaEdgeConfig
+from modeling_llama_edge import LlamaEdgeForCausalLM
+from id_mapper import UnifiedIdMapper
+
+# Load configuration
+config = LlamaEdgeConfig()
+
+# Initialize model
+model = LlamaEdgeForCausalLM(config)
+
+# Load weights
+from safetensors.torch import load_file
+state_dict = load_file("model.safetensors")
+model.load_state_dict(state_dict)
+model.eval()
+```
+
+### Running Inference
+
+Use the `inference.py` script to run a prediction example:
+
+```bash
+python inference.py
+```
+
+## Requirements
+
+Install the dependencies:
+
+```bash
+pip install -r requirement.txt
+```
+
+## Note on Hugging Face Integration
+
+To use with `AutoModel.from_pretrained(..., trust_remote_code=True)`, ensure `config.json` is present (generated from `LlamaEdgeConfig`). Use `register_for_auto_class` if uploading to the Hub.
+
+## How to use from Hugging Face Hub
+
+Users can load this model directly from the Hub without cloning the repository.
+
+### 1. Install Dependencies
+
+```bash
+pip install transformers torch numpy huggingface_hub
+```
+
+### 2. Download and Run Code
+
+```python
+import torch
+from transformers import AutoModel, AutoConfig
+from huggingface_hub import hf_hub_download
+import json
+import sys
+
+# 1. Load Model with trust_remote_code=True
+model_id = "your-username/your-model-name"
+
+# This loads the model and the custom configuration
+model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
+model.eval()
+
+# 2. Load the UnifiedIdMapper
+# The mapper/helper files can be downloaded from the hub repo
+mapper_path = hf_hub_download(repo_id=model_id, filename="unified_id_mapper.json")
+# We also need the python class definition for the mapper. 
+# While it is not automatically imported by AutoModel, you can download it or copy it.
+# If you are running this in a script where you can download files:
+
+id_mapper_scipt = hf_hub_download(repo_id=model_id, filename="id_mapper.py")
+# Import UnifiedIdMapper dynamically or just ensure the file is in path
+import sys
+import os
+sys.path.append(os.path.dirname(id_mapper_scipt))
+from id_mapper import UnifiedIdMapper
+
+# Load mapper
+mapper = UnifiedIdMapper.from_file(mapper_path)
+
+# 3. Example Usage (Manual Inference)
+# Your inference logic here, similar to inference.py
+```

config.json

@@ -0,0 +1,23 @@
+{
+  "architectures": [
+    "LlamaEdgeForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_llama_edge.LlamaEdgeConfig",
+    "AutoModel": "modeling_llama_edge.LlamaEdgeForCausalLM"
+  },
+  "dim": 4096,
+  "dtype": "float32",
+  "ffn_dim_multiplier": 1.3,
+  "intermediate_size": 14336,
+  "max_seq_len": 8192,
+  "model_type": "llama_edge",
+  "multiple_of": 256,
+  "n_heads": 32,
+  "n_kv_heads": 8,
+  "n_layers": 32,
+  "norm_eps": 1e-05,
+  "rope_theta": 500000.0,
+  "transformers_version": "4.57.3",
+  "vocab_size": 9942
+}

configuration_llama_edge.py

@@ -0,0 +1,34 @@
+from transformers import PretrainedConfig
+from typing import Optional
+
+class LlamaEdgeConfig(PretrainedConfig):
+    model_type = "llama_edge"
+
+    def __init__(
+        self,
+        dim: int = 4096,
+        n_layers: int = 32,
+        n_heads: int = 32,
+        n_kv_heads: int = 8,
+        vocab_size: int = 9942,
+        multiple_of: int = 256,
+        ffn_dim_multiplier: Optional[float] = 1.3,
+        norm_eps: float = 1e-5,
+        rope_theta: float = 500000.0,
+        max_seq_len: int = 8192,
+        intermediate_size: int = 14336,
+        **kwargs,
+    ):
+        self.dim = dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.vocab_size = vocab_size
+        self.multiple_of = multiple_of
+        self.ffn_dim_multiplier = ffn_dim_multiplier
+        self.norm_eps = norm_eps
+        self.rope_theta = rope_theta
+        self.max_seq_len = max_seq_len
+        self.intermediate_size = intermediate_size
+
+        super().__init__(**kwargs)

id_mapper.py

@@ -0,0 +1,65 @@
+from typing import Dict, List, Tuple
+import json
+
+class UnifiedIdMapper:
+    def __init__(self, nodes: Dict[int, str], edges: Dict[int, str]) -> None:
+
+        # since all key in JSON are str, convert them to int
+        nodes = {int(k): v for k, v in nodes.items()}
+        edges = {int(k): v for k, v in edges.items()}
+
+        self.nodes = nodes
+        self.edges = edges
+
+        node_mapping = {old_id: new_id for new_id, old_id in enumerate(sorted(self.nodes.keys()))}
+        edge_mapping = {old_id: new_id for new_id, old_id in enumerate(sorted(edges.keys()))}
+        shift = len(nodes)
+
+        self.old_to_new: Dict[int, Tuple[int, bool]] = {
+            **{old_id: (new_id, False) for old_id, new_id in node_mapping.items()},
+            **{old_id: (new_id + shift, True) for old_id, new_id in edge_mapping.items()},
+        }
+        # reverse mapping: new_id -> (old_id, is_edge)
+        self.new_to_old: Dict[int, Tuple[int, bool]] = {
+            new_id: (old_id, is_edge)
+            for old_id, (new_id, is_edge) in self.old_to_new.items()
+        }
+
+        # Label maps
+        self.old_id_to_label: Dict[int, str] = {**nodes, **edges}
+        self.new_id_to_label: Dict[int, str] = {
+            new_id: self.old_id_to_label[old_id] for old_id, (new_id, _) in self.old_to_new.items()
+        }
+
+        self.label_to_old_ids: Dict[str, List[Tuple[int, bool]]] = {}
+        self.label_to_new_ids: Dict[str, List[Tuple[int, bool]]] = {}
+        for old_id, (new_id, is_edge) in self.old_to_new.items():
+            label = self.old_id_to_label.get(old_id)
+            if label is None:
+                continue
+            self.label_to_old_ids.setdefault(label, []).append((old_id, is_edge))
+            self.label_to_new_ids.setdefault(label, []).append((new_id, is_edge))
+
+    @classmethod
+    def from_file(cls, mapper_path: str):
+        with open(mapper_path, "r") as f:
+            data = json.load(f)
+            return cls(data['nodes'], data['edges'])
+
+    def map_old_id(self, old_id: int) -> Tuple[int, bool]:
+        return self.old_to_new[old_id]
+
+    def map_new_id(self, new_id: int) -> Tuple[int, bool]:
+        return self.new_to_old[new_id]
+
+    def label_from_old_id(self, old_id: int) -> str:
+        return self.old_id_to_label[old_id]
+
+    def label_from_new_id(self, new_id: int) -> str:
+        return self.new_id_to_label[new_id]
+
+    def old_ids_from_label(self, label: str) -> List[Tuple[int, bool]]:
+        return self.label_to_old_ids.get(label, [])
+
+    def new_ids_from_label(self, label: str) -> List[Tuple[int, bool]]:
+        return self.label_to_new_ids.get(label, [])

inference.py

@@ -0,0 +1,120 @@
+import torch
+import json
+from safetensors.torch import load_file
+from id_mapper import UnifiedIdMapper
+from modeling_llama_edge import LlamaEdgeForCausalLM
+from configuration_llama_edge import LlamaEdgeConfig
+
+class ModelWrapper:
+    def __init__(self, mapper_path, model_path, device="cuda"):
+        # Load Mapper
+        print(f"Loading mapper from {mapper_path}...")
+        self.mapper = UnifiedIdMapper.from_file(mapper_path)
+
+        # Initialize Empty Model
+        print("Initializing model...")
+        config = LlamaEdgeConfig() # Use defaults or load from file if exists
+        self.model = LlamaEdgeForCausalLM(config)
+
+        # Load Weights
+        print(f"Loading weights from {model_path}...")
+        state_dict = load_file(model_path, device="cpu")
+        self.model.load_state_dict(state_dict)
+
+        # Set device
+        if device == "cuda" and not torch.cuda.is_available():
+            print("CUDA not available, switching to CPU.")
+            self.device = torch.device("cpu")
+        elif device == "mps": # Handle MPS explicitly if requested or available
+             self.device = torch.device("mps")
+        else:
+            self.device = torch.device(device)
+
+        print(f"Moving model to {self.device}...")
+        self.model.to(self.device)
+        self.model.eval()
+
+    def predict(self, old_ids_context):
+        """
+        Args:
+            old_ids_context: List of old IDs defining the context.
+        Returns:
+            sorted_predictions: List of (prob, old_id, label) sorted by probability descending.
+        """
+        # 1. Convert context list of old IDs to new IDs
+        input_ids = []
+        for old_id in old_ids_context:
+            # We assume the input old_ids exist in the mapper
+            new_id, _ = self.mapper.map_old_id(old_id)
+            input_ids.append(new_id)
+
+        # 2. Run inference
+        # Create tensor on result device (batch size = 1)
+        model_input = torch.tensor([input_ids], dtype=torch.long, device=self.device)
+
+        with torch.no_grad():
+            logits = self.model(model_input)
+            # Get logits for the last token in the sequence
+            last_token_logits = logits[0, -1, :]
+            probs = torch.softmax(last_token_logits, dim=-1)
+
+        # 3. Sort by probability descending
+        sorted_probs, sorted_indices = torch.sort(probs, descending=True)
+
+        sorted_probs = sorted_probs.tolist()
+        sorted_indices = sorted_indices.tolist() # These indices are the new_ids
+
+        # 4. Create result list with mapping applied
+        results = []
+        for prob, new_id in zip(sorted_probs, sorted_indices):
+            try:
+                # map_new_id returns (old_id, is_edge)
+                old_id, _ = self.mapper.map_new_id(new_id)
+                label = self.mapper.label_from_new_id(new_id)
+                results.append((prob, old_id, label))
+            except KeyError:
+                # Handle indices not in mapper (e.g., padding tokens)
+                results.append((prob, -1, "<PAD/UNK>"))
+
+        return results
+
+def main():
+    # Define paths
+    mapper_path = "unified_id_mapper.json"
+    model_path = "model.safetensors"
+
+    # Check for device availability
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+    print(f"Using device: {device}")
+
+    # Instantiate the wrapper
+    wrapper = ModelWrapper(mapper_path, model_path, device=device)
+
+    # Input instance take from first line of the training data
+    input_ids = [108, 112, 117, 234, 421, 582, 601, 608, 940, 941, 948, 1008, 1009, 1076, 1094, 1095, 1125, 1188, 1251, 1275, 1365, 1415, 1522, 1687, 1948, 1977, 2025, 47178924, 47185647]
+    target_edge_id = 47182521
+    target_edge_label = "/people/person/place_of_birth"
+
+    predictions = wrapper.predict(input_ids)
+    print(f"Input old IDs: {input_ids}")
+    print(f"Target edge old ID: {target_edge_id}, Label: {target_edge_label}")
+
+    print("Top 10 Predictions:")
+    for rank, (prob, pred_old_id, pred_label) in enumerate(predictions[:10], start=1):
+        print(f"  Rank {rank}: Old ID {pred_old_id}, Label: {pred_label}, Probability: {prob:.6f}")
+
+    # Check if the target is in top 10
+    top_10_old_ids = [pred_old_id for _, pred_old_id, _ in predictions[:10]]
+    if target_edge_id in top_10_old_ids:
+        print(f"Target edge old ID {target_edge_id} found in top 10 predictions.")
+    else:
+        print(f"Target edge old ID {target_edge_id} NOT found in top 10 predictions.")
+    print("-" * 50)
+
+if __name__ == "__main__":
+    main()

model-00001-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd6efe946e1d7ffeccb33ae66d908d10a87c25a2cc93432d677b6750f6ae1a8c
+size 4927788328

model-00002-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71521cd818e18c57340cb0bae75534611f6904c8c774a7bdbc28192b9aab65f8
+size 4999812608

model-00003-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8567a1262dc2f7a4bf40e4b62ad192357dcda0e20c3cff337c86dfc518ba2ccd
+size 4832007088

model-00004-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57bcf7c994d8943add2838392e45875a7ff3477a16f43ac763c2058a00b44071
+size 4999812648

model-00005-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c962e351aa178f4ef5fc5e55a92fd56b5acd73d565aa30695f7962bff7a523e
+size 4999812648

model-00006-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6abc5bfd2dd4c5f6a4d9d7ec83a170c5be89c2865d3f433cc2568f631686d6de
+size 3484929560

model.safetensors.index.json

@@ -0,0 +1,299 @@
+{
+  "metadata": {
+    "total_parameters": 7061032960,
+    "total_size": 28244131840
+  },
+  "weight_map": {
+    "layers.0.attention.wk.weight": "model-00001-of-00006.safetensors",
+    "layers.0.attention.wo.weight": "model-00001-of-00006.safetensors",
+    "layers.0.attention.wq.weight": "model-00001-of-00006.safetensors",
+    "layers.0.attention.wv.weight": "model-00001-of-00006.safetensors",
+    "layers.0.attention_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.0.feed_forward.w1.weight": "model-00001-of-00006.safetensors",
+    "layers.0.feed_forward.w2.weight": "model-00001-of-00006.safetensors",
+    "layers.0.feed_forward.w3.weight": "model-00001-of-00006.safetensors",
+    "layers.0.ffn_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.1.attention.wk.weight": "model-00001-of-00006.safetensors",
+    "layers.1.attention.wo.weight": "model-00001-of-00006.safetensors",
+    "layers.1.attention.wq.weight": "model-00001-of-00006.safetensors",
+    "layers.1.attention.wv.weight": "model-00001-of-00006.safetensors",
+    "layers.1.attention_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.1.feed_forward.w1.weight": "model-00001-of-00006.safetensors",
+    "layers.1.feed_forward.w2.weight": "model-00001-of-00006.safetensors",
+    "layers.1.feed_forward.w3.weight": "model-00001-of-00006.safetensors",
+    "layers.1.ffn_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.10.attention.wk.weight": "model-00002-of-00006.safetensors",
+    "layers.10.attention.wo.weight": "model-00002-of-00006.safetensors",
+    "layers.10.attention.wq.weight": "model-00002-of-00006.safetensors",
+    "layers.10.attention.wv.weight": "model-00002-of-00006.safetensors",
+    "layers.10.attention_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.10.feed_forward.w1.weight": "model-00002-of-00006.safetensors",
+    "layers.10.feed_forward.w2.weight": "model-00002-of-00006.safetensors",
+    "layers.10.feed_forward.w3.weight": "model-00002-of-00006.safetensors",
+    "layers.10.ffn_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.11.attention.wk.weight": "model-00002-of-00006.safetensors",
+    "layers.11.attention.wo.weight": "model-00002-of-00006.safetensors",
+    "layers.11.attention.wq.weight": "model-00002-of-00006.safetensors",
+    "layers.11.attention.wv.weight": "model-00002-of-00006.safetensors",
+    "layers.11.attention_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.11.feed_forward.w1.weight": "model-00003-of-00006.safetensors",
+    "layers.11.feed_forward.w2.weight": "model-00003-of-00006.safetensors",
+    "layers.11.feed_forward.w3.weight": "model-00003-of-00006.safetensors",
+    "layers.11.ffn_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.12.attention.wk.weight": "model-00003-of-00006.safetensors",
+    "layers.12.attention.wo.weight": "model-00003-of-00006.safetensors",
+    "layers.12.attention.wq.weight": "model-00003-of-00006.safetensors",
+    "layers.12.attention.wv.weight": "model-00003-of-00006.safetensors",
+    "layers.12.attention_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.12.feed_forward.w1.weight": "model-00003-of-00006.safetensors",
+    "layers.12.feed_forward.w2.weight": "model-00003-of-00006.safetensors",
+    "layers.12.feed_forward.w3.weight": "model-00003-of-00006.safetensors",
+    "layers.12.ffn_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.13.attention.wk.weight": "model-00003-of-00006.safetensors",
+    "layers.13.attention.wo.weight": "model-00003-of-00006.safetensors",
+    "layers.13.attention.wq.weight": "model-00003-of-00006.safetensors",
+    "layers.13.attention.wv.weight": "model-00003-of-00006.safetensors",
+    "layers.13.attention_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.13.feed_forward.w1.weight": "model-00003-of-00006.safetensors",
+    "layers.13.feed_forward.w2.weight": "model-00003-of-00006.safetensors",
+    "layers.13.feed_forward.w3.weight": "model-00003-of-00006.safetensors",
+    "layers.13.ffn_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.14.attention.wk.weight": "model-00003-of-00006.safetensors",
+    "layers.14.attention.wo.weight": "model-00003-of-00006.safetensors",
+    "layers.14.attention.wq.weight": "model-00003-of-00006.safetensors",
+    "layers.14.attention.wv.weight": "model-00003-of-00006.safetensors",
+    "layers.14.attention_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.14.feed_forward.w1.weight": "model-00003-of-00006.safetensors",
+    "layers.14.feed_forward.w2.weight": "model-00003-of-00006.safetensors",
+    "layers.14.feed_forward.w3.weight": "model-00003-of-00006.safetensors",
+    "layers.14.ffn_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.15.attention.wk.weight": "model-00003-of-00006.safetensors",
+    "layers.15.attention.wo.weight": "model-00003-of-00006.safetensors",
+    "layers.15.attention.wq.weight": "model-00003-of-00006.safetensors",
+    "layers.15.attention.wv.weight": "model-00003-of-00006.safetensors",
+    "layers.15.attention_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.15.feed_forward.w1.weight": "model-00003-of-00006.safetensors",
+    "layers.15.feed_forward.w2.weight": "model-00003-of-00006.safetensors",
+    "layers.15.feed_forward.w3.weight": "model-00003-of-00006.safetensors",
+    "layers.15.ffn_norm.weight": "model-00003-of-00006.safetensors",
+    "layers.16.attention.wk.weight": "model-00003-of-00006.safetensors",
+    "layers.16.attention.wo.weight": "model-00003-of-00006.safetensors",
+    "layers.16.attention.wq.weight": "model-00003-of-00006.safetensors",
+    "layers.16.attention.wv.weight": "model-00003-of-00006.safetensors",
+    "layers.16.attention_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.16.feed_forward.w1.weight": "model-00003-of-00006.safetensors",
+    "layers.16.feed_forward.w2.weight": "model-00003-of-00006.safetensors",
+    "layers.16.feed_forward.w3.weight": "model-00004-of-00006.safetensors",
+    "layers.16.ffn_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.17.attention.wk.weight": "model-00004-of-00006.safetensors",
+    "layers.17.attention.wo.weight": "model-00004-of-00006.safetensors",
+    "layers.17.attention.wq.weight": "model-00004-of-00006.safetensors",
+    "layers.17.attention.wv.weight": "model-00004-of-00006.safetensors",
+    "layers.17.attention_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.17.feed_forward.w1.weight": "model-00004-of-00006.safetensors",
+    "layers.17.feed_forward.w2.weight": "model-00004-of-00006.safetensors",
+    "layers.17.feed_forward.w3.weight": "model-00004-of-00006.safetensors",
+    "layers.17.ffn_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.18.attention.wk.weight": "model-00004-of-00006.safetensors",
+    "layers.18.attention.wo.weight": "model-00004-of-00006.safetensors",
+    "layers.18.attention.wq.weight": "model-00004-of-00006.safetensors",
+    "layers.18.attention.wv.weight": "model-00004-of-00006.safetensors",
+    "layers.18.attention_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.18.feed_forward.w1.weight": "model-00004-of-00006.safetensors",
+    "layers.18.feed_forward.w2.weight": "model-00004-of-00006.safetensors",
+    "layers.18.feed_forward.w3.weight": "model-00004-of-00006.safetensors",
+    "layers.18.ffn_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.19.attention.wk.weight": "model-00004-of-00006.safetensors",
+    "layers.19.attention.wo.weight": "model-00004-of-00006.safetensors",
+    "layers.19.attention.wq.weight": "model-00004-of-00006.safetensors",
+    "layers.19.attention.wv.weight": "model-00004-of-00006.safetensors",
+    "layers.19.attention_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.19.feed_forward.w1.weight": "model-00004-of-00006.safetensors",
+    "layers.19.feed_forward.w2.weight": "model-00004-of-00006.safetensors",
+    "layers.19.feed_forward.w3.weight": "model-00004-of-00006.safetensors",
+    "layers.19.ffn_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.2.attention.wk.weight": "model-00001-of-00006.safetensors",
+    "layers.2.attention.wo.weight": "model-00001-of-00006.safetensors",
+    "layers.2.attention.wq.weight": "model-00001-of-00006.safetensors",
+    "layers.2.attention.wv.weight": "model-00001-of-00006.safetensors",
+    "layers.2.attention_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.2.feed_forward.w1.weight": "model-00001-of-00006.safetensors",
+    "layers.2.feed_forward.w2.weight": "model-00001-of-00006.safetensors",
+    "layers.2.feed_forward.w3.weight": "model-00001-of-00006.safetensors",
+    "layers.2.ffn_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.20.attention.wk.weight": "model-00004-of-00006.safetensors",
+    "layers.20.attention.wo.weight": "model-00004-of-00006.safetensors",
+    "layers.20.attention.wq.weight": "model-00004-of-00006.safetensors",
+    "layers.20.attention.wv.weight": "model-00004-of-00006.safetensors",
+    "layers.20.attention_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.20.feed_forward.w1.weight": "model-00004-of-00006.safetensors",
+    "layers.20.feed_forward.w2.weight": "model-00004-of-00006.safetensors",
+    "layers.20.feed_forward.w3.weight": "model-00004-of-00006.safetensors",
+    "layers.20.ffn_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.21.attention.wk.weight": "model-00004-of-00006.safetensors",
+    "layers.21.attention.wo.weight": "model-00004-of-00006.safetensors",
+    "layers.21.attention.wq.weight": "model-00004-of-00006.safetensors",
+    "layers.21.attention.wv.weight": "model-00004-of-00006.safetensors",
+    "layers.21.attention_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.21.feed_forward.w1.weight": "model-00004-of-00006.safetensors",
+    "layers.21.feed_forward.w2.weight": "model-00004-of-00006.safetensors",
+    "layers.21.feed_forward.w3.weight": "model-00004-of-00006.safetensors",
+    "layers.21.ffn_norm.weight": "model-00004-of-00006.safetensors",
+    "layers.22.attention.wk.weight": "model-00004-of-00006.safetensors",
+    "layers.22.attention.wo.weight": "model-00004-of-00006.safetensors",
+    "layers.22.attention.wq.weight": "model-00004-of-00006.safetensors",
+    "layers.22.attention.wv.weight": "model-00004-of-00006.safetensors",
+    "layers.22.attention_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.22.feed_forward.w1.weight": "model-00004-of-00006.safetensors",
+    "layers.22.feed_forward.w2.weight": "model-00005-of-00006.safetensors",
+    "layers.22.feed_forward.w3.weight": "model-00005-of-00006.safetensors",
+    "layers.22.ffn_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.23.attention.wk.weight": "model-00005-of-00006.safetensors",
+    "layers.23.attention.wo.weight": "model-00005-of-00006.safetensors",
+    "layers.23.attention.wq.weight": "model-00005-of-00006.safetensors",
+    "layers.23.attention.wv.weight": "model-00005-of-00006.safetensors",
+    "layers.23.attention_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.23.feed_forward.w1.weight": "model-00005-of-00006.safetensors",
+    "layers.23.feed_forward.w2.weight": "model-00005-of-00006.safetensors",
+    "layers.23.feed_forward.w3.weight": "model-00005-of-00006.safetensors",
+    "layers.23.ffn_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.24.attention.wk.weight": "model-00005-of-00006.safetensors",
+    "layers.24.attention.wo.weight": "model-00005-of-00006.safetensors",
+    "layers.24.attention.wq.weight": "model-00005-of-00006.safetensors",
+    "layers.24.attention.wv.weight": "model-00005-of-00006.safetensors",
+    "layers.24.attention_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.24.feed_forward.w1.weight": "model-00005-of-00006.safetensors",
+    "layers.24.feed_forward.w2.weight": "model-00005-of-00006.safetensors",
+    "layers.24.feed_forward.w3.weight": "model-00005-of-00006.safetensors",
+    "layers.24.ffn_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.25.attention.wk.weight": "model-00005-of-00006.safetensors",
+    "layers.25.attention.wo.weight": "model-00005-of-00006.safetensors",
+    "layers.25.attention.wq.weight": "model-00005-of-00006.safetensors",
+    "layers.25.attention.wv.weight": "model-00005-of-00006.safetensors",
+    "layers.25.attention_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.25.feed_forward.w1.weight": "model-00005-of-00006.safetensors",
+    "layers.25.feed_forward.w2.weight": "model-00005-of-00006.safetensors",
+    "layers.25.feed_forward.w3.weight": "model-00005-of-00006.safetensors",
+    "layers.25.ffn_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.26.attention.wk.weight": "model-00005-of-00006.safetensors",
+    "layers.26.attention.wo.weight": "model-00005-of-00006.safetensors",
+    "layers.26.attention.wq.weight": "model-00005-of-00006.safetensors",
+    "layers.26.attention.wv.weight": "model-00005-of-00006.safetensors",
+    "layers.26.attention_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.26.feed_forward.w1.weight": "model-00005-of-00006.safetensors",
+    "layers.26.feed_forward.w2.weight": "model-00005-of-00006.safetensors",
+    "layers.26.feed_forward.w3.weight": "model-00005-of-00006.safetensors",
+    "layers.26.ffn_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.27.attention.wk.weight": "model-00005-of-00006.safetensors",
+    "layers.27.attention.wo.weight": "model-00005-of-00006.safetensors",
+    "layers.27.attention.wq.weight": "model-00005-of-00006.safetensors",
+    "layers.27.attention.wv.weight": "model-00005-of-00006.safetensors",
+    "layers.27.attention_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.27.feed_forward.w1.weight": "model-00005-of-00006.safetensors",
+    "layers.27.feed_forward.w2.weight": "model-00005-of-00006.safetensors",
+    "layers.27.feed_forward.w3.weight": "model-00005-of-00006.safetensors",
+    "layers.27.ffn_norm.weight": "model-00005-of-00006.safetensors",
+    "layers.28.attention.wk.weight": "model-00005-of-00006.safetensors",
+    "layers.28.attention.wo.weight": "model-00005-of-00006.safetensors",
+    "layers.28.attention.wq.weight": "model-00005-of-00006.safetensors",
+    "layers.28.attention.wv.weight": "model-00005-of-00006.safetensors",
+    "layers.28.attention_norm.weight": "model-00006-of-00006.safetensors",
+    "layers.28.feed_forward.w1.weight": "model-00006-of-00006.safetensors",
+    "layers.28.feed_forward.w2.weight": "model-00006-of-00006.safetensors",
+    "layers.28.feed_forward.w3.weight": "model-00006-of-00006.safetensors",
+    "layers.28.ffn_norm.weight": "model-00006-of-00006.safetensors",
+    "layers.29.attention.wk.weight": "model-00006-of-00006.safetensors",
+    "layers.29.attention.wo.weight": "model-00006-of-00006.safetensors",
+    "layers.29.attention.wq.weight": "model-00006-of-00006.safetensors",
+    "layers.29.attention.wv.weight": "model-00006-of-00006.safetensors",
+    "layers.29.attention_norm.weight": "model-00006-of-00006.safetensors",
+    "layers.29.feed_forward.w1.weight": "model-00006-of-00006.safetensors",
+    "layers.29.feed_forward.w2.weight": "model-00006-of-00006.safetensors",
+    "layers.29.feed_forward.w3.weight": "model-00006-of-00006.safetensors",
+    "layers.29.ffn_norm.weight": "model-00006-of-00006.safetensors",
+    "layers.3.attention.wk.weight": "model-00001-of-00006.safetensors",
+    "layers.3.attention.wo.weight": "model-00001-of-00006.safetensors",
+    "layers.3.attention.wq.weight": "model-00001-of-00006.safetensors",
+    "layers.3.attention.wv.weight": "model-00001-of-00006.safetensors",
+    "layers.3.attention_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.3.feed_forward.w1.weight": "model-00001-of-00006.safetensors",
+    "layers.3.feed_forward.w2.weight": "model-00001-of-00006.safetensors",
+    "layers.3.feed_forward.w3.weight": "model-00001-of-00006.safetensors",
+    "layers.3.ffn_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.30.attention.wk.weight": "model-00006-of-00006.safetensors",
+    "layers.30.attention.wo.weight": "model-00006-of-00006.safetensors",
+    "layers.30.attention.wq.weight": "model-00006-of-00006.safetensors",
+    "layers.30.attention.wv.weight": "model-00006-of-00006.safetensors",
+    "layers.30.attention_norm.weight": "model-00006-of-00006.safetensors",
+    "layers.30.feed_forward.w1.weight": "model-00006-of-00006.safetensors",
+    "layers.30.feed_forward.w2.weight": "model-00006-of-00006.safetensors",
+    "layers.30.feed_forward.w3.weight": "model-00006-of-00006.safetensors",
+    "layers.30.ffn_norm.weight": "model-00006-of-00006.safetensors",
+    "layers.31.attention.wk.weight": "model-00006-of-00006.safetensors",
+    "layers.31.attention.wo.weight": "model-00006-of-00006.safetensors",
+    "layers.31.attention.wq.weight": "model-00006-of-00006.safetensors",
+    "layers.31.attention.wv.weight": "model-00006-of-00006.safetensors",
+    "layers.31.attention_norm.weight": "model-00006-of-00006.safetensors",
+    "layers.31.feed_forward.w1.weight": "model-00006-of-00006.safetensors",
+    "layers.31.feed_forward.w2.weight": "model-00006-of-00006.safetensors",
+    "layers.31.feed_forward.w3.weight": "model-00006-of-00006.safetensors",
+    "layers.31.ffn_norm.weight": "model-00006-of-00006.safetensors",
+    "layers.4.attention.wk.weight": "model-00001-of-00006.safetensors",
+    "layers.4.attention.wo.weight": "model-00001-of-00006.safetensors",
+    "layers.4.attention.wq.weight": "model-00001-of-00006.safetensors",
+    "layers.4.attention.wv.weight": "model-00001-of-00006.safetensors",
+    "layers.4.attention_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.4.feed_forward.w1.weight": "model-00001-of-00006.safetensors",
+    "layers.4.feed_forward.w2.weight": "model-00001-of-00006.safetensors",
+    "layers.4.feed_forward.w3.weight": "model-00001-of-00006.safetensors",
+    "layers.4.ffn_norm.weight": "model-00001-of-00006.safetensors",
+    "layers.5.attention.wk.weight": "model-00001-of-00006.safetensors",
+    "layers.5.attention.wo.weight": "model-00001-of-00006.safetensors",
+    "layers.5.attention.wq.weight": "model-00001-of-00006.safetensors",
+    "layers.5.attention.wv.weight": "model-00001-of-00006.safetensors",
+    "layers.5.attention_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.5.feed_forward.w1.weight": "model-00001-of-00006.safetensors",
+    "layers.5.feed_forward.w2.weight": "model-00002-of-00006.safetensors",
+    "layers.5.feed_forward.w3.weight": "model-00002-of-00006.safetensors",
+    "layers.5.ffn_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.6.attention.wk.weight": "model-00002-of-00006.safetensors",
+    "layers.6.attention.wo.weight": "model-00002-of-00006.safetensors",
+    "layers.6.attention.wq.weight": "model-00002-of-00006.safetensors",
+    "layers.6.attention.wv.weight": "model-00002-of-00006.safetensors",
+    "layers.6.attention_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.6.feed_forward.w1.weight": "model-00002-of-00006.safetensors",
+    "layers.6.feed_forward.w2.weight": "model-00002-of-00006.safetensors",
+    "layers.6.feed_forward.w3.weight": "model-00002-of-00006.safetensors",
+    "layers.6.ffn_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.7.attention.wk.weight": "model-00002-of-00006.safetensors",
+    "layers.7.attention.wo.weight": "model-00002-of-00006.safetensors",
+    "layers.7.attention.wq.weight": "model-00002-of-00006.safetensors",
+    "layers.7.attention.wv.weight": "model-00002-of-00006.safetensors",
+    "layers.7.attention_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.7.feed_forward.w1.weight": "model-00002-of-00006.safetensors",
+    "layers.7.feed_forward.w2.weight": "model-00002-of-00006.safetensors",
+    "layers.7.feed_forward.w3.weight": "model-00002-of-00006.safetensors",
+    "layers.7.ffn_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.8.attention.wk.weight": "model-00002-of-00006.safetensors",
+    "layers.8.attention.wo.weight": "model-00002-of-00006.safetensors",
+    "layers.8.attention.wq.weight": "model-00002-of-00006.safetensors",
+    "layers.8.attention.wv.weight": "model-00002-of-00006.safetensors",
+    "layers.8.attention_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.8.feed_forward.w1.weight": "model-00002-of-00006.safetensors",
+    "layers.8.feed_forward.w2.weight": "model-00002-of-00006.safetensors",
+    "layers.8.feed_forward.w3.weight": "model-00002-of-00006.safetensors",
+    "layers.8.ffn_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.9.attention.wk.weight": "model-00002-of-00006.safetensors",
+    "layers.9.attention.wo.weight": "model-00002-of-00006.safetensors",
+    "layers.9.attention.wq.weight": "model-00002-of-00006.safetensors",
+    "layers.9.attention.wv.weight": "model-00002-of-00006.safetensors",
+    "layers.9.attention_norm.weight": "model-00002-of-00006.safetensors",
+    "layers.9.feed_forward.w1.weight": "model-00002-of-00006.safetensors",
+    "layers.9.feed_forward.w2.weight": "model-00002-of-00006.safetensors",
+    "layers.9.feed_forward.w3.weight": "model-00002-of-00006.safetensors",
+    "layers.9.ffn_norm.weight": "model-00002-of-00006.safetensors",
+    "norm.weight": "model-00006-of-00006.safetensors",
+    "output.weight": "model-00006-of-00006.safetensors",
+    "token_embedding.weight": "model-00001-of-00006.safetensors"
+  }
+}

modeling_llama_edge.py

@@ -0,0 +1,144 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+from transformers import PreTrainedModel
+from configuration_llama_edge import LlamaEdgeConfig
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return output * self.weight
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    # Precompute complex exponentials for Rotary Positional Embeddings (RoPE)
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, multiple_of: int, ffn_dim_multiplier: Optional[float]):
+        super().__init__()
+
+        # If the config provides a specific hidden_dim (intermediate_size), use it directly.
+        # Otherwise, calculate it using the standard Llama formula.
+        if hidden_dim is None:
+            hidden_dim = 4 * dim
+            hidden_dim = int(2 * hidden_dim / 3)
+            if ffn_dim_multiplier is not None:
+                hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        # In Llama 3 8B, this will now be 14336
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+class Attention(nn.Module):
+    def __init__(self, config: LlamaEdgeConfig):
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.n_kv_heads = config.n_kv_heads
+        self.head_dim = config.dim // config.n_heads
+        self.n_rep = self.n_heads // self.n_kv_heads
+
+        self.wq = nn.Linear(config.dim, config.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(config.dim, config.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(config.dim, config.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(config.n_heads * self.head_dim, config.dim, bias=False)
+
+    def forward(self, x, freqs_cis, mask=None):
+        bsz, seqlen, _ = x.shape
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+
+        # Reshape for multi-head attention
+        xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_kv_heads, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_kv_heads, self.head_dim)
+
+        # Apply RoPE
+        # xq, xk = apply_rotary_emb(xq, xk, freqs_cis) 
+        
+        # Repeat K and V heads for GQA (if n_kv_heads < n_heads)
+        if self.n_rep > 1:
+            xk = xk.unsqueeze(3).repeat(1, 1, 1, self.n_rep, 1).reshape(bsz, seqlen, self.n_heads, self.head_dim)
+            xv = xv.unsqueeze(3).repeat(1, 1, 1, self.n_rep, 1).reshape(bsz, seqlen, self.n_heads, self.head_dim)
+
+        # Transpose for attention calculation: (bsz, heads, seqlen, dim)
+        xq = xq.transpose(1, 2)
+        xk = xk.transpose(1, 2)
+        xv = xv.transpose(1, 2)
+
+        # Scaled Dot-Product Attention
+        scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
+        # if mask is not None:
+        #     scores = scores + mask  # Apply causal mask
+
+        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+        output = torch.matmul(scores, xv)
+
+        # Reshape back
+        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        return self.wo(output)
+
+class TransformerBlock(nn.Module):
+    def __init__(self, layer_id: int, config: LlamaEdgeConfig):
+        super().__init__()
+        self.attention = Attention(config)
+        self.feed_forward = FeedForward(
+            dim=config.dim,
+            hidden_dim=config.intermediate_size,
+            multiple_of=config.multiple_of,
+            ffn_dim_multiplier=config.ffn_dim_multiplier,
+        )
+        self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+
+    def forward(self, x, freqs_cis, mask=None):
+        h = x + self.attention(self.attention_norm(x), freqs_cis, mask)
+        out = h + self.feed_forward(self.ffn_norm(h))
+        return out
+
+class LlamaEdgeForCausalLM(PreTrainedModel):
+    config_class = LlamaEdgeConfig
+
+    def __init__(self, config: LlamaEdgeConfig):
+        super().__init__(config)
+        self.token_embedding = nn.Embedding(config.vocab_size, config.dim)
+        self.layers = nn.ModuleList([TransformerBlock(i, config) for i in range(config.n_layers)])
+        self.norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
+
+        # Precompute RoPE frequencies
+        self.freqs_cis = precompute_freqs_cis(
+            config.dim // config.n_heads, config.max_seq_len, config.rope_theta,
+        )
+        
+
+    def forward(self, x):
+        bsz, seqlen = x.shape
+        freqs_cis = self.freqs_cis[:seqlen].to(x.device)
+
+        # Create causal mask
+        mask = torch.full((seqlen, seqlen), float("-inf"), device=x.device)
+        mask = torch.triu(mask, diagonal=1)
+
+        h = self.token_embedding(x)
+
+        for layer in self.layers:
+            h = layer(h, freqs_cis, mask)
+
+        h = self.norm(h)
+        logits = self.output(h)
+        return logits