Upload 8 files
Browse filesInitial release of Alpha Transformer.
Added model weights, configuration files, vocabularies, tokenizer, and inference code.
51M parameter English-French Transformer trained from scratch on OPUS-100.
- README.md +255 -1
- config.json +33 -0
- inference.py +124 -0
- main.py +191 -0
- model.safetensors +3 -0
- src_vocab.json +0 -0
- tgt_vocab.json +0 -0
- tokenizer.py +95 -0
README.md
CHANGED
|
@@ -1,3 +1,257 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Alpha Transformer
|
| 2 |
+
|
| 3 |
+
A custom **Encoder-Decoder Transformer model trained from scratch** for English → French translation.
|
| 4 |
+
|
| 5 |
+
Alpha is a 51M parameter Transformer implemented entirely in PyTorch, including:
|
| 6 |
+
- Multi-Head Attention
|
| 7 |
+
- Rotary Positional Embeddings (RoPE)
|
| 8 |
+
- Encoder-Decoder architecture
|
| 9 |
+
- Cross Attention
|
| 10 |
+
- Custom tokenizer
|
| 11 |
+
- Mixed precision training
|
| 12 |
+
- Checkpointing system
|
| 13 |
+
|
| 14 |
+
The model was trained on the **OPUS-100 English-French translation dataset**.
|
| 15 |
+
|
| 16 |
---
|
| 17 |
+
|
| 18 |
+
## Model Overview
|
| 19 |
+
|
| 20 |
+
| Feature | Value |
|
| 21 |
+
|---|---|
|
| 22 |
+
| Architecture | Encoder-Decoder Transformer |
|
| 23 |
+
| Parameters | 51.26M |
|
| 24 |
+
| Framework | PyTorch |
|
| 25 |
+
| Task | English → French Translation |
|
| 26 |
+
| Dataset | OPUS-100 |
|
| 27 |
+
| Training Samples | ~1M sentences |
|
| 28 |
+
| Embedding Dimension | 256 |
|
| 29 |
+
| Attention Heads | 8 |
|
| 30 |
+
| Feed Forward Dimension | 1024 |
|
| 31 |
+
| Encoder Layers | 3 |
|
| 32 |
+
| Decoder Layers | 3 |
|
| 33 |
+
| Maximum Sequence Length | 64 tokens |
|
| 34 |
+
| Positional Encoding | RoPE |
|
| 35 |
+
|
| 36 |
---
|
| 37 |
+
|
| 38 |
+
# Performance
|
| 39 |
+
|
| 40 |
+
Evaluation on the OPUS-100 test set:
|
| 41 |
+
|
| 42 |
+
| Metric | Result |
|
| 43 |
+
|---|---|
|
| 44 |
+
| Token Accuracy | 62.8% |
|
| 45 |
+
|
| 46 |
+
The model was trained from scratch without using pretrained weights.
|
| 47 |
+
|
| 48 |
+
---
|
| 49 |
+
|
| 50 |
+
# Architecture
|
| 51 |
+
|
| 52 |
+
## Encoder
|
| 53 |
+
|
| 54 |
+
Each encoder block contains:
|
| 55 |
+
|
| 56 |
+
Input Embeddings
|
| 57 |
+
|
|
| 58 |
+
v
|
| 59 |
+
LayerNorm
|
| 60 |
+
|
|
| 61 |
+
Multi-Head Self Attention + RoPE
|
| 62 |
+
|
|
| 63 |
+
Residual Connection
|
| 64 |
+
|
|
| 65 |
+
LayerNorm
|
| 66 |
+
|
|
| 67 |
+
Feed Forward Network
|
| 68 |
+
|
|
| 69 |
+
Residual Connection
|
| 70 |
+
|
| 71 |
+
## Decoder
|
| 72 |
+
|
| 73 |
+
Each decoder block contains:
|
| 74 |
+
|
| 75 |
+
Target Embeddings
|
| 76 |
+
|
|
| 77 |
+
LayerNorm
|
| 78 |
+
|
|
| 79 |
+
Masked Self Attention + RoPE
|
| 80 |
+
|
|
| 81 |
+
Residual Connection
|
| 82 |
+
|
|
| 83 |
+
LayerNorm
|
| 84 |
+
|
|
| 85 |
+
Cross Attention
|
| 86 |
+
|
|
| 87 |
+
Residual Connection
|
| 88 |
+
|
|
| 89 |
+
LayerNorm
|
| 90 |
+
|
|
| 91 |
+
Feed Forward Network
|
| 92 |
+
|
|
| 93 |
+
Residual Connection
|
| 94 |
+
|
| 95 |
+
---
|
| 96 |
+
|
| 97 |
+
# Files
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
Alpha/
|
| 101 |
+
│
|
| 102 |
+
├── config.json
|
| 103 |
+
│ └── Model architecture configuration
|
| 104 |
+
│
|
| 105 |
+
├── model.safetensors
|
| 106 |
+
│ └── Trained model weights
|
| 107 |
+
│
|
| 108 |
+
├── src_vocab.json
|
| 109 |
+
│ └── English vocabulary
|
| 110 |
+
│
|
| 111 |
+
├── tgt_vocab.json
|
| 112 |
+
│ └── French vocabulary
|
| 113 |
+
│
|
| 114 |
+
├── tokenizer.py
|
| 115 |
+
│ └── Custom tokenizer implementation
|
| 116 |
+
│
|
| 117 |
+
├── main.py
|
| 118 |
+
│ └── Transformer architecture
|
| 119 |
+
│
|
| 120 |
+
├── inference.py
|
| 121 |
+
│ └── Translation script
|
| 122 |
+
│
|
| 123 |
+
└── README.md
|
| 124 |
+
|
| 125 |
+
---
|
| 126 |
+
|
| 127 |
+
# Installation
|
| 128 |
+
|
| 129 |
+
Clone the repository:
|
| 130 |
+
|
| 131 |
+
```bash
|
| 132 |
+
git clone https://github.com/MehdiSkilll24/Alpha.git
|
| 133 |
+
cd Alpha
|
| 134 |
+
|
| 135 |
+
Install dependencies:
|
| 136 |
+
|
| 137 |
+
pip install torch safetensors
|
| 138 |
+
|
| 139 |
+
Usage
|
| 140 |
+
|
| 141 |
+
Run inference:
|
| 142 |
+
|
| 143 |
+
python inference.py
|
| 144 |
+
|
| 145 |
+
Example:
|
| 146 |
+
|
| 147 |
+
Input:
|
| 148 |
+
> hello, how are you?
|
| 149 |
+
|
| 150 |
+
Output:
|
| 151 |
+
> bonjour, comment allez-vous ?
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
Loading the Model Manually
|
| 156 |
+
import torch
|
| 157 |
+
import json
|
| 158 |
+
from safetensors.torch import load_file
|
| 159 |
+
from main import Transformer
|
| 160 |
+
|
| 161 |
+
with open("config.json") as f:
|
| 162 |
+
config = json.load(f)
|
| 163 |
+
|
| 164 |
+
weights = load_file("model.safetensors")
|
| 165 |
+
|
| 166 |
+
model = Transformer(
|
| 167 |
+
src_vocab_size=config["src_vocab_size"],
|
| 168 |
+
tgt_vocab_size=config["tgt_vocab_size"],
|
| 169 |
+
d_model=config["d_model"],
|
| 170 |
+
num_heads=config["num_heads"],
|
| 171 |
+
d_ff=config["d_ff"]
|
| 172 |
+
)
|
| 173 |
+
|
| 174 |
+
model.load_state_dict(weights)
|
| 175 |
+
model.eval()
|
| 176 |
+
|
| 177 |
+
|
| 178 |
+
Training
|
| 179 |
+
|
| 180 |
+
The model was trained using:
|
| 181 |
+
|
| 182 |
+
AdamW optimizer
|
| 183 |
+
Mixed precision training (FP16)
|
| 184 |
+
ReduceLROnPlateau scheduler
|
| 185 |
+
Custom PyTorch training loop
|
| 186 |
+
|
| 187 |
+
Training configuration:
|
| 188 |
+
|
| 189 |
+
Learning Rate: 3e-4
|
| 190 |
+
Batch Size: 32
|
| 191 |
+
Epochs: 18
|
| 192 |
+
Dataset
|
| 193 |
+
|
| 194 |
+
Training data:
|
| 195 |
+
|
| 196 |
+
OPUS-100 English-French translation dataset.
|
| 197 |
+
|
| 198 |
+
The dataset contains parallel English/French sentence pairs collected from various sources.
|
| 199 |
+
|
| 200 |
+
Limitations
|
| 201 |
+
|
| 202 |
+
This model is designed as an educational and experimental Transformer implementation.
|
| 203 |
+
|
| 204 |
+
Limitations:
|
| 205 |
+
|
| 206 |
+
Vocabulary is based on a custom word-level tokenizer
|
| 207 |
+
Limited context length
|
| 208 |
+
No pretrained initialization
|
| 209 |
+
Smaller dataset compared to production translation systems
|
| 210 |
+
Token accuracy does not directly represent BLEU score
|
| 211 |
+
Future Improvements
|
| 212 |
+
|
| 213 |
+
Possible improvements:
|
| 214 |
+
|
| 215 |
+
Replace word tokenizer with BPE/SentencePiece
|
| 216 |
+
Increase model depth
|
| 217 |
+
Train on larger datasets
|
| 218 |
+
Add beam search decoding
|
| 219 |
+
Evaluate with BLEU / COMET
|
| 220 |
+
Implement KV-cache for faster inference
|
| 221 |
+
Motivation
|
| 222 |
+
|
| 223 |
+
Alpha was built to understand and implement the Transformer architecture from first principles.
|
| 224 |
+
|
| 225 |
+
Rather than using existing libraries, the project implements the core components manually:
|
| 226 |
+
|
| 227 |
+
Attention calculations
|
| 228 |
+
Positional embeddings
|
| 229 |
+
Encoder/decoder blocks
|
| 230 |
+
Training pipeline
|
| 231 |
+
Model serialization
|
| 232 |
+
|
| 233 |
+
The goal is to better understand how modern language models are built internally.
|
| 234 |
+
|
| 235 |
+
Source Code
|
| 236 |
+
|
| 237 |
+
GitHub:
|
| 238 |
+
|
| 239 |
+
https://github.com/MehdiSkilll24/Alpha
|
| 240 |
+
|
| 241 |
+
Model weights:
|
| 242 |
+
|
| 243 |
+
https://huggingface.co/MehdiSkilll24/Alpha
|
| 244 |
+
|
| 245 |
+
License
|
| 246 |
+
|
| 247 |
+
MIT License
|
| 248 |
+
|
| 249 |
+
|
| 250 |
+
One suggestion before you upload: add **3-5 real translation examples** from your model. That will make the HF page look 10x more alive. A model card with "62.8% accuracy" is okay; a model card with:
|
| 251 |
+
|
| 252 |
+
|
| 253 |
+
English:
|
| 254 |
+
"The weather is beautiful today."
|
| 255 |
+
|
| 256 |
+
Alpha:
|
| 257 |
+
"Le temps est magnifique aujourd'hui."
|
config.json
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name": "Alpha-Transformer",
|
| 3 |
+
"model_type": "encoder-decoder-transformer",
|
| 4 |
+
|
| 5 |
+
"d_model": 256,
|
| 6 |
+
"num_heads": 8,
|
| 7 |
+
"d_ff": 1024,
|
| 8 |
+
|
| 9 |
+
"num_encoder_layers": 3,
|
| 10 |
+
"num_decoder_layers": 3,
|
| 11 |
+
|
| 12 |
+
"src_vocab_size": 51278,
|
| 13 |
+
"tgt_vocab_size": 63558,
|
| 14 |
+
|
| 15 |
+
"max_seq_length": 64,
|
| 16 |
+
"rope_max_seq_length": 1024,
|
| 17 |
+
|
| 18 |
+
"positional_encoding": "RoPE",
|
| 19 |
+
"activation": "ReLU",
|
| 20 |
+
"normalization": "Pre-LayerNorm",
|
| 21 |
+
|
| 22 |
+
"num_parameters": 51262022,
|
| 23 |
+
|
| 24 |
+
"dataset": "OPUS-100",
|
| 25 |
+
"task": "English-French machine translation",
|
| 26 |
+
|
| 27 |
+
"framework": "PyTorch",
|
| 28 |
+
|
| 29 |
+
"pad_token_id": 0,
|
| 30 |
+
"unk_token_id": 1,
|
| 31 |
+
"bos_token_id": 2,
|
| 32 |
+
"eos_token_id": 3
|
| 33 |
+
}
|
inference.py
ADDED
|
@@ -0,0 +1,124 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from main import Transformer
|
| 3 |
+
from tokenizer import Tokenizer
|
| 4 |
+
import json
|
| 5 |
+
from safetensors.torch import load_file
|
| 6 |
+
|
| 7 |
+
with open("config.json") as f:
|
| 8 |
+
config = json.load(f)
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def rebuild_tokenizer(word_to_idx):
|
| 12 |
+
tok = Tokenizer()
|
| 13 |
+
|
| 14 |
+
tok.word_to_idx = word_to_idx
|
| 15 |
+
|
| 16 |
+
tok.idx_to_word = {
|
| 17 |
+
int(idx): word
|
| 18 |
+
for word, idx in word_to_idx.items()
|
| 19 |
+
}
|
| 20 |
+
|
| 21 |
+
return tok
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def translate(model, src_tokenizer, tgt_tokenizer, sentence, device, max_len=50):
|
| 25 |
+
model.eval()
|
| 26 |
+
|
| 27 |
+
src_ids = src_tokenizer.encode(sentence)
|
| 28 |
+
src_tensor = torch.tensor([src_ids]).to(device)
|
| 29 |
+
|
| 30 |
+
decoder_input = torch.tensor(
|
| 31 |
+
[[tgt_tokenizer.word_to_idx["<bos>"]]]
|
| 32 |
+
).to(device)
|
| 33 |
+
|
| 34 |
+
with torch.no_grad():
|
| 35 |
+
|
| 36 |
+
for _ in range(max_len):
|
| 37 |
+
|
| 38 |
+
logits = model(src_tensor, decoder_input)
|
| 39 |
+
|
| 40 |
+
next_token = logits[:, -1, :].argmax(dim=-1).item()
|
| 41 |
+
|
| 42 |
+
decoder_input = torch.cat(
|
| 43 |
+
[
|
| 44 |
+
decoder_input,
|
| 45 |
+
torch.tensor([[next_token]]).to(device)
|
| 46 |
+
],
|
| 47 |
+
dim=1
|
| 48 |
+
)
|
| 49 |
+
|
| 50 |
+
if next_token == tgt_tokenizer.word_to_idx["<eos>"]:
|
| 51 |
+
break
|
| 52 |
+
|
| 53 |
+
output_ids = decoder_input[0].tolist()
|
| 54 |
+
|
| 55 |
+
return tgt_tokenizer.decode(output_ids)
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def main():
|
| 59 |
+
|
| 60 |
+
device = torch.device(
|
| 61 |
+
"cuda" if torch.cuda.is_available() else "cpu"
|
| 62 |
+
)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
# Load vocabularies
|
| 66 |
+
with open("src_vocab.json", "r", encoding="utf-8") as f:
|
| 67 |
+
src_vocab = json.load(f)
|
| 68 |
+
|
| 69 |
+
with open("tgt_vocab.json", "r", encoding="utf-8") as f:
|
| 70 |
+
tgt_vocab = json.load(f)
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
# Rebuild tokenizers
|
| 74 |
+
src_tokenizer = rebuild_tokenizer(src_vocab)
|
| 75 |
+
tgt_tokenizer = rebuild_tokenizer(tgt_vocab)
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
# Load config
|
| 79 |
+
with open("config.json", "r", encoding="utf-8") as f:
|
| 80 |
+
config = json.load(f)
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
# Build model
|
| 84 |
+
model = Transformer(
|
| 85 |
+
src_vocab_size=config["src_vocab_size"],
|
| 86 |
+
tgt_vocab_size=config["tgt_vocab_size"],
|
| 87 |
+
d_model=config["d_model"],
|
| 88 |
+
num_heads=config["num_heads"],
|
| 89 |
+
d_ff=config["d_ff"]
|
| 90 |
+
).to(device)
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
# Load weights
|
| 94 |
+
weights = load_file(
|
| 95 |
+
"model.safetensors"
|
| 96 |
+
)
|
| 97 |
+
|
| 98 |
+
model.load_state_dict(weights)
|
| 99 |
+
|
| 100 |
+
model.eval()
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
print("Type an English sentence (or 'quit' to exit):")
|
| 104 |
+
|
| 105 |
+
while True:
|
| 106 |
+
|
| 107 |
+
sentence = input("> ")
|
| 108 |
+
|
| 109 |
+
if sentence.lower() == "quit":
|
| 110 |
+
break
|
| 111 |
+
|
| 112 |
+
translation = translate(
|
| 113 |
+
model,
|
| 114 |
+
src_tokenizer,
|
| 115 |
+
tgt_tokenizer,
|
| 116 |
+
sentence,
|
| 117 |
+
device
|
| 118 |
+
)
|
| 119 |
+
|
| 120 |
+
print(f"Translation: {translation}")
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
if __name__ == "__main__":
|
| 124 |
+
main()
|
main.py
ADDED
|
@@ -0,0 +1,191 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch.nn as nn
|
| 2 |
+
import torch
|
| 3 |
+
|
| 4 |
+
class RoPE(nn.Module):
|
| 5 |
+
def __init__(self, dim, max_seq_len=1024):
|
| 6 |
+
super().__init__()
|
| 7 |
+
|
| 8 |
+
theta = 10000
|
| 9 |
+
|
| 10 |
+
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
|
| 11 |
+
|
| 12 |
+
positions = torch.arange(max_seq_len)
|
| 13 |
+
|
| 14 |
+
angles = torch.outer(positions, freqs)
|
| 15 |
+
|
| 16 |
+
self.register_buffer(
|
| 17 |
+
"cos",
|
| 18 |
+
angles.cos()
|
| 19 |
+
)
|
| 20 |
+
|
| 21 |
+
self.register_buffer(
|
| 22 |
+
"sin",
|
| 23 |
+
angles.sin()
|
| 24 |
+
)
|
| 25 |
+
|
| 26 |
+
def rotate_half(self, x):
|
| 27 |
+
x1 = x[..., ::2]
|
| 28 |
+
x2 = x[..., 1::2]
|
| 29 |
+
|
| 30 |
+
return torch.stack(
|
| 31 |
+
(-x2,x1),
|
| 32 |
+
dim=-1
|
| 33 |
+
).flatten(-2)
|
| 34 |
+
|
| 35 |
+
def forward(self, x):
|
| 36 |
+
|
| 37 |
+
seq_len = x.shape[-2]
|
| 38 |
+
cos = self.cos[:seq_len]
|
| 39 |
+
sin = self.sin[:seq_len]
|
| 40 |
+
|
| 41 |
+
cos = torch.repeat_interleave(cos, 2, dim=-1)
|
| 42 |
+
sin = torch.repeat_interleave(sin, 2, dim=-1)
|
| 43 |
+
|
| 44 |
+
return x*cos + self.rotate_half(x)*sin
|
| 45 |
+
|
| 46 |
+
class MHA(nn.Module):
|
| 47 |
+
def __init__(self, d_model, num_heads):
|
| 48 |
+
super().__init__()
|
| 49 |
+
super().__init__()
|
| 50 |
+
|
| 51 |
+
self.num_heads = num_heads
|
| 52 |
+
self.d_model = d_model
|
| 53 |
+
|
| 54 |
+
self.Q = nn.Linear(d_model, d_model)
|
| 55 |
+
self.K = nn.Linear(d_model, d_model)
|
| 56 |
+
self.V = nn.Linear(d_model, d_model)
|
| 57 |
+
self.W_o = nn.Linear(d_model, d_model)
|
| 58 |
+
|
| 59 |
+
self.d_k = d_model // num_heads
|
| 60 |
+
self.rope = RoPE(self.d_k)
|
| 61 |
+
|
| 62 |
+
def forward(self, x, kv_input=None, mask=None):
|
| 63 |
+
batch_size = x.size(0)
|
| 64 |
+
# x.shape = (2, 3, 64) 64 dictates that there will be 64 Q, K and V PER TOKEN
|
| 65 |
+
|
| 66 |
+
if kv_input is None:
|
| 67 |
+
kv_input = x
|
| 68 |
+
|
| 69 |
+
Q = self.Q(x) # we do q_1 = XWq = 'W_1'*'e1' + ... + 'W_64'*'e64' W_i being the weight at i ... until we comp q_64, if no kv_input, it means we're doing encoding
|
| 70 |
+
K = self.K(kv_input) # ,, ,,
|
| 71 |
+
V = self.V(kv_input) # ,, ,,
|
| 72 |
+
# We do this to ensure all Q,K and V's matrices have access to the input embeddings
|
| 73 |
+
|
| 74 |
+
# Manual length extraction to avoid mismatch in size from encoder/decoder
|
| 75 |
+
q_seq_length = x.shape[1]
|
| 76 |
+
kv_seq_length = kv_input.shape[1]
|
| 77 |
+
|
| 78 |
+
# We then split all this data through different heads so each one goes for a raneg window for each token (split ranges -> specialized heads)
|
| 79 |
+
Q = Q.view(batch_size, q_seq_length, self.num_heads, self.d_k)
|
| 80 |
+
K = K.view(batch_size, kv_seq_length, self.num_heads, self.d_k)
|
| 81 |
+
V = V.view(batch_size, kv_seq_length, self.num_heads, self.d_k)
|
| 82 |
+
|
| 83 |
+
# So we get (batch, heads, seq_len, d_k) last 2 dims are the ones that get multiplied by Pytorch
|
| 84 |
+
# The reason we put seq_len, d_k is because we gotta multiply them indep. of how many heads or batches, what matters is the consistent output per head per batch
|
| 85 |
+
# Also, since we're computing relationships between tokens, there must be seq_len and d_k, so we can compare all tokens with all tokens for every head in every batch
|
| 86 |
+
Q = Q.transpose(1,2)
|
| 87 |
+
K = K.transpose(1,2)
|
| 88 |
+
V = V.transpose(1,2)
|
| 89 |
+
|
| 90 |
+
Q = self.rope(Q)
|
| 91 |
+
K = self.rope(K)
|
| 92 |
+
|
| 93 |
+
scores = Q @ K.transpose(-2, -1) / self.d_k**0.5 #we swap d_k with seq_len so we get #seq_len * #seq_len matrix as a result
|
| 94 |
+
|
| 95 |
+
if mask is not None:
|
| 96 |
+
scores += mask
|
| 97 |
+
|
| 98 |
+
attn_weights = torch.softmax(scores, dim=-1)
|
| 99 |
+
output = attn_weights @ V # contiguous vector of shape (batch, heads, seq, d_k)
|
| 100 |
+
output = output.transpose(1,2).contiguous().view(batch_size, q_seq_length, self.d_model) # transposing makes it breaks contiguousy, so we add .contiguous() and then .view() to reshape it with concatenation
|
| 101 |
+
output = self.W_o(output) # expressing the new shape correctly accross heads
|
| 102 |
+
|
| 103 |
+
return output
|
| 104 |
+
|
| 105 |
+
class MLP(nn.Module):
|
| 106 |
+
|
| 107 |
+
def __init__(self, d_model, d_ff):
|
| 108 |
+
super().__init__()
|
| 109 |
+
self.fc1 = nn.Linear(d_model, d_ff)
|
| 110 |
+
self.relu = nn.ReLU()
|
| 111 |
+
self.fc2 = nn.Linear(d_ff, d_model)
|
| 112 |
+
|
| 113 |
+
def forward(self, x):
|
| 114 |
+
x = self.fc1(x)
|
| 115 |
+
x = self.relu(x)
|
| 116 |
+
x = self.fc2(x)
|
| 117 |
+
|
| 118 |
+
return x
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
class EncoderBlock(nn.Module):
|
| 122 |
+
def __init__(self, d_model, num_heads, d_ff):
|
| 123 |
+
super().__init__()
|
| 124 |
+
self.mha = MHA(d_model, num_heads)
|
| 125 |
+
self.norm1 = nn.LayerNorm(d_model)
|
| 126 |
+
self.norm2 = nn.LayerNorm(d_model)
|
| 127 |
+
self.mlp = MLP(d_model, d_ff)
|
| 128 |
+
|
| 129 |
+
def forward(self, x):
|
| 130 |
+
normed = self.norm1(x)
|
| 131 |
+
attn_out = self.mha(normed)
|
| 132 |
+
x = x + attn_out
|
| 133 |
+
|
| 134 |
+
normed2 = self.norm2(x)
|
| 135 |
+
mlp_out = self.mlp(normed2)
|
| 136 |
+
x = x + mlp_out
|
| 137 |
+
|
| 138 |
+
return x
|
| 139 |
+
|
| 140 |
+
class DecoderBlock(nn.Module):
|
| 141 |
+
def __init__(self, d_model, num_heads, d_ff):
|
| 142 |
+
super().__init__()
|
| 143 |
+
self.self_attn = MHA(d_model, num_heads)
|
| 144 |
+
self.cross_attn = MHA(d_model, num_heads)
|
| 145 |
+
self.norm1 = nn.LayerNorm(d_model)
|
| 146 |
+
self.norm2 = nn.LayerNorm(d_model)
|
| 147 |
+
self.norm3 = nn.LayerNorm(d_model)
|
| 148 |
+
self.mlp = MLP(d_model, d_ff)
|
| 149 |
+
|
| 150 |
+
def forward(self, x, encoder_output):
|
| 151 |
+
causal_mask = torch.triu(torch.full((x.shape[1], x.shape[1]), float('-inf'), device=x.device), diagonal=1)
|
| 152 |
+
normed = self.norm1(x)
|
| 153 |
+
attn_out = self.self_attn(normed, mask = causal_mask)
|
| 154 |
+
x = x + attn_out
|
| 155 |
+
|
| 156 |
+
normed2 = self.norm2(x)
|
| 157 |
+
cross_out = self.cross_attn(normed2, encoder_output)
|
| 158 |
+
x = x + cross_out
|
| 159 |
+
|
| 160 |
+
normed3 = self.norm3(x)
|
| 161 |
+
mlp_out = self.mlp(normed3)
|
| 162 |
+
x = x + mlp_out
|
| 163 |
+
|
| 164 |
+
return x
|
| 165 |
+
|
| 166 |
+
class Transformer(nn.Module):
|
| 167 |
+
def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, d_ff):
|
| 168 |
+
super().__init__()
|
| 169 |
+
|
| 170 |
+
self.src_embedding = nn.Embedding(src_vocab_size, d_model)
|
| 171 |
+
self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
|
| 172 |
+
self.encoder = nn.ModuleList([EncoderBlock(d_model, num_heads, d_ff) for _ in range(3)])
|
| 173 |
+
self.decoder = nn.ModuleList([DecoderBlock(d_model, num_heads, d_ff) for _ in range(3)])
|
| 174 |
+
self.output_proj = nn.Linear(d_model, tgt_vocab_size)
|
| 175 |
+
|
| 176 |
+
def forward(self, src_tokens, tgt_tokens):
|
| 177 |
+
src = self.src_embedding(src_tokens)
|
| 178 |
+
tgt = self.tgt_embedding(tgt_tokens)
|
| 179 |
+
|
| 180 |
+
for block in self.encoder:
|
| 181 |
+
src = block(src)
|
| 182 |
+
|
| 183 |
+
encoder_output = src
|
| 184 |
+
|
| 185 |
+
for block in self.decoder:
|
| 186 |
+
tgt = block(tgt, encoder_output)
|
| 187 |
+
|
| 188 |
+
logits = self.output_proj(tgt)
|
| 189 |
+
return logits
|
| 190 |
+
|
| 191 |
+
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:4533ddea34e9a72f8cd42c11d2933ed5cd0363dbc4f67c9adf9a226afa779cce
|
| 3 |
+
size 206241552
|
src_vocab.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
tgt_vocab.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
tokenizer.py
ADDED
|
@@ -0,0 +1,95 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import re
|
| 2 |
+
from collections import Counter
|
| 3 |
+
|
| 4 |
+
class Tokenizer:
|
| 5 |
+
def __init__(self, min_freq = 5):
|
| 6 |
+
self.min_freq = min_freq
|
| 7 |
+
|
| 8 |
+
self.word_to_idx = {}
|
| 9 |
+
self.idx_to_word = {}
|
| 10 |
+
|
| 11 |
+
self.special_tokens = {
|
| 12 |
+
"<pad>" : 0,
|
| 13 |
+
"<unk>" : 1,
|
| 14 |
+
"<bos>" : 2,
|
| 15 |
+
"<eos>" : 3,
|
| 16 |
+
}
|
| 17 |
+
|
| 18 |
+
def __len__(self):
|
| 19 |
+
return len(self.word_to_idx)
|
| 20 |
+
|
| 21 |
+
def tokenize(self, text):
|
| 22 |
+
|
| 23 |
+
# raw text -> list of words
|
| 24 |
+
|
| 25 |
+
text = text.lower()
|
| 26 |
+
|
| 27 |
+
# split words and punctuations
|
| 28 |
+
|
| 29 |
+
tokens = re.findall(r"\w+|[^\w\s]", text)
|
| 30 |
+
|
| 31 |
+
return tokens
|
| 32 |
+
|
| 33 |
+
def build_vocab(self, texts):
|
| 34 |
+
|
| 35 |
+
# create word -> id mapping
|
| 36 |
+
|
| 37 |
+
counter = Counter()
|
| 38 |
+
|
| 39 |
+
for text in texts:
|
| 40 |
+
tokens = self.tokenize(text)
|
| 41 |
+
counter.update(tokens)
|
| 42 |
+
|
| 43 |
+
self.word_to_idx = self.special_tokens.copy()
|
| 44 |
+
|
| 45 |
+
for word, freq in sorted(counter.items()):
|
| 46 |
+
if freq >= self.min_freq:
|
| 47 |
+
if word not in self.word_to_idx:
|
| 48 |
+
self.word_to_idx[word] = len(self.word_to_idx)
|
| 49 |
+
|
| 50 |
+
self.idx_to_word = {
|
| 51 |
+
idx: word for word, idx in self.word_to_idx.items()
|
| 52 |
+
}
|
| 53 |
+
|
| 54 |
+
def encode(self, text):
|
| 55 |
+
|
| 56 |
+
# text -> ids
|
| 57 |
+
text = text.lower()
|
| 58 |
+
|
| 59 |
+
tokens = self.tokenize(text)
|
| 60 |
+
|
| 61 |
+
ids = [
|
| 62 |
+
self.word_to_idx.get(
|
| 63 |
+
token,
|
| 64 |
+
self.word_to_idx["<unk>"]
|
| 65 |
+
)
|
| 66 |
+
for token in tokens
|
| 67 |
+
]
|
| 68 |
+
|
| 69 |
+
ids = (
|
| 70 |
+
[self.word_to_idx["<bos>"]]
|
| 71 |
+
+ ids
|
| 72 |
+
+[self.word_to_idx["<eos>"]]
|
| 73 |
+
)
|
| 74 |
+
|
| 75 |
+
return ids
|
| 76 |
+
|
| 77 |
+
def decode(self, ids):
|
| 78 |
+
words = []
|
| 79 |
+
|
| 80 |
+
for idx in ids:
|
| 81 |
+
|
| 82 |
+
word = self.idx_to_word.get(
|
| 83 |
+
idx,
|
| 84 |
+
"<unk>"
|
| 85 |
+
)
|
| 86 |
+
|
| 87 |
+
if word not in [
|
| 88 |
+
"<pad>",
|
| 89 |
+
"<bos>",
|
| 90 |
+
"<eos>",
|
| 91 |
+
"<unk>"
|
| 92 |
+
]:
|
| 93 |
+
words.append(word)
|
| 94 |
+
|
| 95 |
+
return " ".join(words)
|