Spaces:
Build error
Build error
Commit
·
3aa6cf7
1
Parent(s):
878f4f2
huggingface app
Browse files- .gitignore +67 -0
- README.md +24 -8
- app.py +115 -0
- requirements.txt +4 -0
- src/config/__init__.py +1 -0
- src/config/model_config.py +9 -0
- src/models/__init__.py +4 -0
- src/models/attention.py +36 -0
- src/models/block.py +16 -0
- src/models/gpt.py +110 -0
- src/models/mlp.py +15 -0
- src/utils/device_utils.py +40 -0
.gitignore
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Python
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.py[cod]
|
| 4 |
+
*$py.class
|
| 5 |
+
*.so
|
| 6 |
+
.Python
|
| 7 |
+
build/
|
| 8 |
+
develop-eggs/
|
| 9 |
+
dist/
|
| 10 |
+
downloads/
|
| 11 |
+
eggs/
|
| 12 |
+
.eggs/
|
| 13 |
+
lib/
|
| 14 |
+
lib64/
|
| 15 |
+
parts/
|
| 16 |
+
sdist/
|
| 17 |
+
var/
|
| 18 |
+
wheels/
|
| 19 |
+
*.egg-info/
|
| 20 |
+
.installed.cfg
|
| 21 |
+
*.egg
|
| 22 |
+
|
| 23 |
+
# Virtual Environment
|
| 24 |
+
venv/
|
| 25 |
+
env/
|
| 26 |
+
ENV/
|
| 27 |
+
.env
|
| 28 |
+
.venv
|
| 29 |
+
|
| 30 |
+
# IDE
|
| 31 |
+
.idea/
|
| 32 |
+
.vscode/
|
| 33 |
+
*.swp
|
| 34 |
+
*.swo
|
| 35 |
+
.DS_Store
|
| 36 |
+
|
| 37 |
+
# Project specific
|
| 38 |
+
input.txt
|
| 39 |
+
*.pt
|
| 40 |
+
*.pth
|
| 41 |
+
wandb/
|
| 42 |
+
logs/
|
| 43 |
+
checkpoints/
|
| 44 |
+
|
| 45 |
+
# Jupyter Notebook
|
| 46 |
+
.ipynb_checkpoints
|
| 47 |
+
*.ipynb
|
| 48 |
+
|
| 49 |
+
# Distribution / packaging
|
| 50 |
+
.Python
|
| 51 |
+
*.pyc
|
| 52 |
+
|
| 53 |
+
# Unit test / coverage reports
|
| 54 |
+
htmlcov/
|
| 55 |
+
.tox/
|
| 56 |
+
.coverage
|
| 57 |
+
.coverage.*
|
| 58 |
+
.cache
|
| 59 |
+
nosetests.xml
|
| 60 |
+
coverage.xml
|
| 61 |
+
*.cover
|
| 62 |
+
.hypothesis/
|
| 63 |
+
|
| 64 |
+
# mypy
|
| 65 |
+
.mypy_cache/
|
| 66 |
+
.dmypy.json
|
| 67 |
+
dmypy.json
|
README.md
CHANGED
|
@@ -1,14 +1,30 @@
|
|
| 1 |
---
|
| 2 |
-
title:
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
sdk: streamlit
|
| 7 |
-
sdk_version: 1.
|
| 8 |
app_file: app.py
|
| 9 |
pinned: false
|
| 10 |
-
license: apache-2.0
|
| 11 |
-
short_description: gpt-text-generator
|
| 12 |
---
|
| 13 |
|
| 14 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
+
title: GPT Text Generator
|
| 3 |
+
emoji: 🤖
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: red
|
| 6 |
sdk: streamlit
|
| 7 |
+
sdk_version: 1.25.0
|
| 8 |
app_file: app.py
|
| 9 |
pinned: false
|
|
|
|
|
|
|
| 10 |
---
|
| 11 |
|
| 12 |
+
# GPT Text Generator
|
| 13 |
+
|
| 14 |
+
A Streamlit app that generates text using a custom GPT model.
|
| 15 |
+
|
| 16 |
+
## Features
|
| 17 |
+
- Text generation from prompts
|
| 18 |
+
- Adjustable generation length
|
| 19 |
+
- Multiple sequence generation
|
| 20 |
+
|
| 21 |
+
## Usage
|
| 22 |
+
1. Enter your prompt text
|
| 23 |
+
2. Adjust the additional tokens to predict
|
| 24 |
+
3. Select number of sequences
|
| 25 |
+
4. Click Generate
|
| 26 |
+
|
| 27 |
+
## Model Details
|
| 28 |
+
- Architecture: GPT
|
| 29 |
+
- Vocabulary: GPT-2 tokenizer
|
| 30 |
+
- Training Data: Custom dataset
|
app.py
ADDED
|
@@ -0,0 +1,115 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import streamlit as st
|
| 2 |
+
import torch
|
| 3 |
+
import tiktoken
|
| 4 |
+
import sys
|
| 5 |
+
import os
|
| 6 |
+
import logging
|
| 7 |
+
import warnings
|
| 8 |
+
|
| 9 |
+
# Configure logging and warnings
|
| 10 |
+
logging.getLogger('streamlit').setLevel(logging.ERROR)
|
| 11 |
+
warnings.filterwarnings('ignore', message='.*torch.classes.*')
|
| 12 |
+
warnings.filterwarnings('ignore', category=FutureWarning)
|
| 13 |
+
|
| 14 |
+
# Add the project root to Python path
|
| 15 |
+
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 16 |
+
|
| 17 |
+
from src.config.model_config import GPTConfig
|
| 18 |
+
from src.models.gpt import GPT
|
| 19 |
+
from src.utils.device_utils import get_device
|
| 20 |
+
|
| 21 |
+
@st.cache_resource
|
| 22 |
+
def load_model():
|
| 23 |
+
device = get_device()
|
| 24 |
+
config = GPTConfig()
|
| 25 |
+
model = GPT(config)
|
| 26 |
+
|
| 27 |
+
# Load the trained weights
|
| 28 |
+
checkpoint = torch.load('checkpoints/final_model.pt', map_location=device, weights_only=True)
|
| 29 |
+
|
| 30 |
+
# Handle pruned weights
|
| 31 |
+
state_dict = checkpoint['model_state_dict']
|
| 32 |
+
new_state_dict = {}
|
| 33 |
+
|
| 34 |
+
for key in model.state_dict().keys():
|
| 35 |
+
if key.endswith('.weight'):
|
| 36 |
+
# Check if this is a pruned weight
|
| 37 |
+
orig_key = key[:-7] + '.weight_orig' if key.endswith('.weight') else key
|
| 38 |
+
mask_key = key[:-7] + '.weight_mask' if key.endswith('.weight') else key
|
| 39 |
+
|
| 40 |
+
if orig_key in state_dict and mask_key in state_dict:
|
| 41 |
+
# Reconstruct the pruned weight
|
| 42 |
+
new_state_dict[key] = state_dict[orig_key] * state_dict[mask_key]
|
| 43 |
+
else:
|
| 44 |
+
# Use the weight as is
|
| 45 |
+
new_state_dict[key] = state_dict[key] if key in state_dict else model.state_dict()[key]
|
| 46 |
+
else:
|
| 47 |
+
# Copy non-weight parameters as is
|
| 48 |
+
new_state_dict[key] = state_dict[key] if key in state_dict else model.state_dict()[key]
|
| 49 |
+
|
| 50 |
+
# Load the processed state dict
|
| 51 |
+
model.load_state_dict(new_state_dict)
|
| 52 |
+
|
| 53 |
+
# Convert back to float32 for inference
|
| 54 |
+
model = model.float()
|
| 55 |
+
model.to(device)
|
| 56 |
+
model.eval()
|
| 57 |
+
|
| 58 |
+
return model, device
|
| 59 |
+
|
| 60 |
+
def generate_text(model, prompt, max_length=100, num_return_sequences=1, device='cpu'):
|
| 61 |
+
tokenizer = tiktoken.get_encoding('gpt2')
|
| 62 |
+
input_tokens = tokenizer.encode(prompt)
|
| 63 |
+
x = torch.tensor(input_tokens).unsqueeze(0).repeat(num_return_sequences, 1)
|
| 64 |
+
x = x.to(device)
|
| 65 |
+
|
| 66 |
+
# Calculate final length (input length + requested additional tokens)
|
| 67 |
+
input_length = x.size(1)
|
| 68 |
+
target_length = input_length + max_length
|
| 69 |
+
|
| 70 |
+
# Generate text
|
| 71 |
+
with torch.no_grad():
|
| 72 |
+
while x.size(1) < target_length:
|
| 73 |
+
logits = model(x)[0]
|
| 74 |
+
next_token_logits = logits[:, -1, :]
|
| 75 |
+
probs = torch.softmax(next_token_logits, dim=-1)
|
| 76 |
+
next_token = torch.multinomial(probs, num_samples=1)
|
| 77 |
+
x = torch.cat((x, next_token), dim=1)
|
| 78 |
+
|
| 79 |
+
# Print token information once before generating sequences
|
| 80 |
+
st.text(f"Size of Input tokens: {input_length}, Additional tokens to be predicted: {max_length}, Total tokens to be generated: {x.size(1)}")
|
| 81 |
+
|
| 82 |
+
# Decode generated sequences
|
| 83 |
+
generated_texts = []
|
| 84 |
+
for i in range(num_return_sequences):
|
| 85 |
+
tokens = x[i].tolist()
|
| 86 |
+
text = tokenizer.decode(tokens)
|
| 87 |
+
generated_texts.append(text)
|
| 88 |
+
|
| 89 |
+
return generated_texts
|
| 90 |
+
|
| 91 |
+
# Streamlit UI
|
| 92 |
+
st.title("GPT Text Generator")
|
| 93 |
+
|
| 94 |
+
# Load model
|
| 95 |
+
model, device = load_model()
|
| 96 |
+
|
| 97 |
+
# Input form
|
| 98 |
+
prompt = st.text_area("Enter your prompt:", "Once upon a time")
|
| 99 |
+
max_length = st.slider("Predict additional text of length:", min_value=10, max_value=100, value=10)
|
| 100 |
+
num_sequences = st.slider("Number of sequences to generate:", 1, 5, 1)
|
| 101 |
+
|
| 102 |
+
if st.button("Generate"):
|
| 103 |
+
with st.spinner("Generating text..."):
|
| 104 |
+
generated_texts = generate_text(
|
| 105 |
+
model=model,
|
| 106 |
+
prompt=prompt,
|
| 107 |
+
max_length=max_length,
|
| 108 |
+
num_return_sequences=num_sequences,
|
| 109 |
+
device=device
|
| 110 |
+
)
|
| 111 |
+
|
| 112 |
+
# Display results
|
| 113 |
+
for i, text in enumerate(generated_texts, 1):
|
| 114 |
+
st.write(f"\nSequence {i}:")
|
| 115 |
+
st.write(text)
|
requirements.txt
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
streamlit
|
| 2 |
+
torch>=1.13.0 # For quantization support
|
| 3 |
+
tiktoken
|
| 4 |
+
transformers
|
src/config/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
from .model_config import GPTConfig
|
src/config/model_config.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from dataclasses import dataclass
|
| 2 |
+
|
| 3 |
+
@dataclass
|
| 4 |
+
class GPTConfig:
|
| 5 |
+
block_size: int = 1024 # max sequence length
|
| 6 |
+
vocab_size: int = 50257 # number of tokens: 50,000 BPE merges + 256 bytes tokens + 1 <|endoftext|> token
|
| 7 |
+
n_layer: int = 12 # number of layers
|
| 8 |
+
n_head: int = 12 # number of heads
|
| 9 |
+
n_embd: int = 768 # embedding dimension
|
src/models/__init__.py
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from .gpt import GPT
|
| 2 |
+
from .block import Block
|
| 3 |
+
from .attention import CausalSelfAttention
|
| 4 |
+
from .mlp import MLP
|
src/models/attention.py
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import math
|
| 2 |
+
import torch
|
| 3 |
+
import torch.nn as nn
|
| 4 |
+
from torch.nn import functional as F
|
| 5 |
+
|
| 6 |
+
class CausalSelfAttention(nn.Module):
|
| 7 |
+
def __init__(self, config):
|
| 8 |
+
super().__init__()
|
| 9 |
+
assert config.n_embd % config.n_head == 0
|
| 10 |
+
# key, query, value projections for all heads, but in a batch
|
| 11 |
+
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
|
| 12 |
+
# output projection
|
| 13 |
+
self.c_proj = nn.Linear(config.n_embd, config.n_embd)
|
| 14 |
+
self.c_proj.NANGPT_SCALE_INIT = 1
|
| 15 |
+
# regularization
|
| 16 |
+
self.n_head = config.n_head
|
| 17 |
+
self.n_embd = config.n_embd
|
| 18 |
+
self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
|
| 19 |
+
.view(1, 1, config.block_size, config.block_size))
|
| 20 |
+
|
| 21 |
+
def forward(self, x):
|
| 22 |
+
B, T, C = x.size()
|
| 23 |
+
qkv = self.c_attn(x)
|
| 24 |
+
q, k, v = qkv.split(self.n_embd, dim=2)
|
| 25 |
+
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
|
| 26 |
+
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
|
| 27 |
+
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
|
| 28 |
+
|
| 29 |
+
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
|
| 30 |
+
att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
|
| 31 |
+
att = F.softmax(att, dim=-1)
|
| 32 |
+
y = att @ v
|
| 33 |
+
|
| 34 |
+
y = y.transpose(1, 2).contiguous().view(B, T, C)
|
| 35 |
+
y = self.c_proj(y)
|
| 36 |
+
return y
|
src/models/block.py
ADDED
|
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch.nn as nn
|
| 2 |
+
from .attention import CausalSelfAttention
|
| 3 |
+
from .mlp import MLP
|
| 4 |
+
|
| 5 |
+
class Block(nn.Module):
|
| 6 |
+
def __init__(self, config):
|
| 7 |
+
super().__init__()
|
| 8 |
+
self.ln_1 = nn.LayerNorm(config.n_embd)
|
| 9 |
+
self.attn = CausalSelfAttention(config)
|
| 10 |
+
self.ln_2 = nn.LayerNorm(config.n_embd)
|
| 11 |
+
self.mlp = MLP(config)
|
| 12 |
+
|
| 13 |
+
def forward(self, x):
|
| 14 |
+
x = x + self.attn(self.ln_1(x))
|
| 15 |
+
x = x + self.mlp(self.ln_2(x))
|
| 16 |
+
return x
|
src/models/gpt.py
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
import torch.nn as nn
|
| 3 |
+
from torch.nn import functional as F
|
| 4 |
+
from .block import Block
|
| 5 |
+
import sys
|
| 6 |
+
import os
|
| 7 |
+
|
| 8 |
+
# Add the project root to Python path
|
| 9 |
+
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
| 10 |
+
|
| 11 |
+
from src.config.model_config import GPTConfig
|
| 12 |
+
|
| 13 |
+
class GPT(nn.Module):
|
| 14 |
+
def __init__(self, config):
|
| 15 |
+
super().__init__()
|
| 16 |
+
self.config = config
|
| 17 |
+
|
| 18 |
+
self.transformer = nn.ModuleDict(dict(
|
| 19 |
+
wte = nn.Embedding(config.vocab_size, config.n_embd),
|
| 20 |
+
wpe = nn.Embedding(config.block_size, config.n_embd),
|
| 21 |
+
h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
|
| 22 |
+
ln_f = nn.LayerNorm(config.n_embd),
|
| 23 |
+
))
|
| 24 |
+
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
|
| 25 |
+
|
| 26 |
+
# weight sharing
|
| 27 |
+
self.transformer.wte.weight = self.lm_head.weight
|
| 28 |
+
|
| 29 |
+
# weight initialization
|
| 30 |
+
self.apply(self._init_weights)
|
| 31 |
+
|
| 32 |
+
def _init_weights(self, module):
|
| 33 |
+
if isinstance(module, nn.Linear):
|
| 34 |
+
std = 0.02
|
| 35 |
+
if hasattr(module, 'NANGPT_SCALE_INIT'):
|
| 36 |
+
std *= (2 * self.config.n_layer) ** -0.5
|
| 37 |
+
torch.nn.init.normal_(module.weight, mean=0.0, std=std)
|
| 38 |
+
if module.bias is not None:
|
| 39 |
+
torch.nn.init.zeros_(module.bias)
|
| 40 |
+
elif isinstance(module, nn.Embedding):
|
| 41 |
+
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
| 42 |
+
|
| 43 |
+
def forward(self, idx, targets=None):
|
| 44 |
+
B, T = idx.size()
|
| 45 |
+
assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, block size is only {self.config.block_size}"
|
| 46 |
+
|
| 47 |
+
pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
|
| 48 |
+
pos_emb = self.transformer.wpe(pos)
|
| 49 |
+
tok_emb = self.transformer.wte(idx)
|
| 50 |
+
x = tok_emb + pos_emb
|
| 51 |
+
|
| 52 |
+
for block in self.transformer.h:
|
| 53 |
+
x = block(x)
|
| 54 |
+
|
| 55 |
+
x = self.transformer.ln_f(x)
|
| 56 |
+
logits = self.lm_head(x)
|
| 57 |
+
|
| 58 |
+
loss = None
|
| 59 |
+
if targets is not None:
|
| 60 |
+
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
|
| 61 |
+
return logits, loss
|
| 62 |
+
|
| 63 |
+
@classmethod
|
| 64 |
+
def from_pretrained(cls, model_type):
|
| 65 |
+
"""Loads pretrained GPT-2 model weights from huggingface"""
|
| 66 |
+
assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
|
| 67 |
+
from transformers import GPT2LMHeadModel
|
| 68 |
+
print("loading weights from pretrained gpt: %s" % model_type)
|
| 69 |
+
|
| 70 |
+
# n_layer, n_head and n_embd are determined from model_type
|
| 71 |
+
config_args = {
|
| 72 |
+
'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params
|
| 73 |
+
'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
|
| 74 |
+
'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
|
| 75 |
+
'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
|
| 76 |
+
}[model_type]
|
| 77 |
+
config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
|
| 78 |
+
config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
|
| 79 |
+
|
| 80 |
+
# create a from-scratch initialized minGPT model
|
| 81 |
+
config = GPTConfig(**config_args)
|
| 82 |
+
model = GPT(config)
|
| 83 |
+
sd = model.state_dict()
|
| 84 |
+
sd_keys = sd.keys()
|
| 85 |
+
sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param
|
| 86 |
+
|
| 87 |
+
# init a huggingface/transformers model
|
| 88 |
+
model_hf = GPT2LMHeadModel.from_pretrained(model_type)
|
| 89 |
+
sd_hf = model_hf.state_dict()
|
| 90 |
+
|
| 91 |
+
# copy while ensuring all of the parameters are aligned and match in names and shapes
|
| 92 |
+
sd_keys_hf = sd_hf.keys()
|
| 93 |
+
sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
|
| 94 |
+
sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
|
| 95 |
+
transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
|
| 96 |
+
|
| 97 |
+
assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
|
| 98 |
+
for k in sd_keys_hf:
|
| 99 |
+
if any(k.endswith(w) for w in transposed):
|
| 100 |
+
# special treatment for the Conv1D weights we need to transpose
|
| 101 |
+
assert sd_hf[k].shape[::-1] == sd[k].shape
|
| 102 |
+
with torch.no_grad():
|
| 103 |
+
sd[k].copy_(sd_hf[k].t())
|
| 104 |
+
else:
|
| 105 |
+
# vanilla copy over the other parameters
|
| 106 |
+
assert sd_hf[k].shape == sd[k].shape
|
| 107 |
+
with torch.no_grad():
|
| 108 |
+
sd[k].copy_(sd_hf[k])
|
| 109 |
+
|
| 110 |
+
return model
|
src/models/mlp.py
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch.nn as nn
|
| 2 |
+
|
| 3 |
+
class MLP(nn.Module):
|
| 4 |
+
def __init__(self, config):
|
| 5 |
+
super().__init__()
|
| 6 |
+
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
|
| 7 |
+
self.gelu = nn.GELU(approximate='tanh')
|
| 8 |
+
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
|
| 9 |
+
self.c_proj.NANOGPT_SCALE_INIT = 1
|
| 10 |
+
|
| 11 |
+
def forward(self, x):
|
| 12 |
+
x = self.c_fc(x)
|
| 13 |
+
x = self.gelu(x)
|
| 14 |
+
x = self.c_proj(x)
|
| 15 |
+
return x
|
src/utils/device_utils.py
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
import os
|
| 3 |
+
|
| 4 |
+
def get_device():
|
| 5 |
+
if torch.cuda.is_available():
|
| 6 |
+
return 'cuda'
|
| 7 |
+
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
| 8 |
+
return "mps"
|
| 9 |
+
return 'cpu'
|
| 10 |
+
|
| 11 |
+
def set_seed(seed=1337):
|
| 12 |
+
torch.manual_seed(seed)
|
| 13 |
+
if torch.cuda.is_available():
|
| 14 |
+
torch.cuda.manual_seed(seed)
|
| 15 |
+
|
| 16 |
+
def save_model(model, optimizer, loss, epoch, path='checkpoints/model.pt'):
|
| 17 |
+
os.makedirs('checkpoints', exist_ok=True)
|
| 18 |
+
|
| 19 |
+
# Convert model to half precision
|
| 20 |
+
model_to_save = model.half()
|
| 21 |
+
|
| 22 |
+
# Save in half precision
|
| 23 |
+
torch.save({
|
| 24 |
+
'model_state_dict': model_to_save.state_dict(),
|
| 25 |
+
'loss': loss,
|
| 26 |
+
'epoch': epoch
|
| 27 |
+
}, path, _use_new_zipfile_serialization=False)
|
| 28 |
+
|
| 29 |
+
# Convert back to original precision
|
| 30 |
+
model.float()
|
| 31 |
+
print(f"Model saved to {path}")
|
| 32 |
+
|
| 33 |
+
def load_model(model, optimizer=None, path='checkpoints/model.pt'):
|
| 34 |
+
if os.path.exists(path):
|
| 35 |
+
checkpoint = torch.load(path, weights_only=True)
|
| 36 |
+
model.load_state_dict(checkpoint['model_state_dict'])
|
| 37 |
+
if optimizer is not None and 'optimizer_state_dict' in checkpoint:
|
| 38 |
+
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
| 39 |
+
return checkpoint.get('epoch', 0), checkpoint.get('loss', None)
|
| 40 |
+
return 0, None
|