Upload 14 files

app.py

@@ -0,0 +1,244 @@
+import os
+import pickle
+import torch
+import random
+import subprocess
+import re
+import pretty_midi
+import gradio as gr
+from contextlib import nullcontext
+from model import GPTConfig, GPT
+from pedalboard import Pedalboard, Reverb, Compressor, Gain, Limiter
+from pedalboard.io import AudioFile
+
+in_space = os.getenv("SYSTEM") == "spaces"
+
+temp_dir = 'temp'
+os.makedirs(temp_dir, exist_ok=True)
+
+init_from = 'resume'
+out_dir = 'checkpoints'
+ckpt_load = 'model.pt'
+
+start = "000000000000\n"
+num_samples = 1
+max_new_tokens = 564
+
+seed = random.randint(1, 100000)
+torch.manual_seed(seed)
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16'
+compile = False
+exec(open('configurator.py').read())
+
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+device_type = 'cpu' if 'cuda' in device else 'cpu'
+ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
+ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+
+if init_from == 'resume':
+    ckpt_path = os.path.join(out_dir, ckpt_load)
+    checkpoint = torch.load(ckpt_path, map_location=device, weights_only=True)
+    gptconf = GPTConfig(**checkpoint['model_args'])
+    model = GPT(gptconf)
+    state_dict = checkpoint['model']
+    unwanted_prefix = '_orig_mod.'
+    for k, v in list(state_dict.items()):
+        if k.startswith(unwanted_prefix):
+            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+    model.load_state_dict(state_dict)
+elif init_from.startswith('gpt2'):
+    model = GPT.from_pretrained(init_from, dict(dropout=0.0))
+
+model.eval()
+model.to(device)
+if compile:
+    model = torch.compile(model)
+
+tokenizer = re.compile(r'000000000000|\d{1}|\n')
+
+meta_path = os.path.join('', 'meta.pkl')
+with open(meta_path, 'rb') as f:
+    meta = pickle.load(f)
+    stoi = meta.get('stoi', None)
+    itos = meta.get('itos', None)
+
+def encode(text):
+    matches = tokenizer.findall(text)
+    return [stoi[c] for c in matches]
+
+def decode(encoded):
+    return ''.join([itos[i] for i in encoded])
+
+def clear_midi(dir):
+    for file in os.listdir(dir):
+        if file.endswith('.mid'):
+            os.remove(os.path.join(dir, file))
+
+clear_midi(temp_dir)
+
+
+def generate_midi(temperature, top_k):
+    start_ids = encode(start)
+    x = (torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...])
+        
+    midi_events = []
+    seq_count = 0
+
+    with torch.no_grad():
+        for _ in range(num_samples):
+            sequence = []
+            y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
+            tkn_seq = decode(y[0].tolist())
+            lines = tkn_seq.splitlines()
+
+            for event in lines:
+                if event.startswith(start.strip()):
+                    if sequence:
+                        midi_events.append(sequence)
+                        sequence = []
+                    seq_count += 1
+                elif event.strip() == "":
+                    continue
+                else:
+                    try:
+                        p = int(event[0:2])
+                        v = int(event[2:4])
+                        s = int(event[4:8])
+                        e = int(event[8:12])
+                    except ValueError:
+                        p, v, s, e = 0, 0, 0, 0
+                    sequence.append({'file_name': f'nanompc_{seq_count:02d}', 'pitch': p, 'velocity': v, 'start': s, 'end': e})
+
+            if sequence:
+                midi_events.append(sequence)
+
+    round_bars = []
+    
+    for sequence in midi_events:
+        filtered_sequence = []
+        for event in sequence:
+            if event['start'] < 768 and event['end'] <= 768:
+                filtered_sequence.append(event)
+        if filtered_sequence:
+            round_bars.append(filtered_sequence)
+
+    midi_events = round_bars
+
+    return midi_events
+
+
+def write_midi(midi_events, bpm):
+    midi_data = pretty_midi.PrettyMIDI(initial_tempo=bpm, resolution=96)
+    midi_data.time_signature_changes.append(pretty_midi.containers.TimeSignature(4, 4, 0))
+    instrument = pretty_midi.Instrument(0)
+    midi_data.instruments.append(instrument)
+
+    for sequence in midi_events:
+        for event in sequence:
+            pitch = event['pitch']
+            velocity = event['velocity']
+            start = midi_data.tick_to_time(event['start'])
+            end = midi_data.tick_to_time(event['end'])
+            note = pretty_midi.Note(pitch=pitch, velocity=velocity, start=start, end=end)
+            instrument.notes.append(note)
+
+    midi_path = os.path.join(temp_dir, 'output.mid')
+    midi_data.write(midi_path)
+    print(f"Generated: {midi_path}")
+
+
+def render_wav(midi_file):
+    
+    sf2_dir = 'sf2_kits'
+    audio_format = 's16'
+    sample_rate = '44100'
+    gain = '2.0'
+
+    sf2_files = [f for f in os.listdir(sf2_dir) if f.endswith('.sf2')]
+    if not sf2_files:
+        raise ValueError("No SoundFont (.sf2) file found in directory.")
+
+    sf2_file = os.path.join(sf2_dir, random.choice(sf2_files))
+    print(sf2_file)
+    output_wav = os.path.join(temp_dir, 'output.wav')
+
+    with open(os.devnull, 'w') as devnull:
+        command = [
+            'fluidsynth', '-ni', sf2_file, midi_file, '-F', output_wav, '-r', str(sample_rate), 
+            '-o', f'audio.file.format={audio_format}', '-g', str(gain)
+        ]
+        subprocess.call(command, stdout=devnull, stderr=devnull)
+
+    return output_wav
+
+
+def render_sfx(wav_raw, settings):
+    wav_fx = wav_raw
+
+    for setting in settings:
+        board = setting['board']
+
+        with AudioFile(wav_raw) as f:
+            with AudioFile(wav_fx, 'w', f.samplerate, f.num_channels) as o:
+                while f.tell() < f.frames:
+                    chunk = f.read(int(f.samplerate))
+                    effected = board(chunk, f.samplerate, reset=False)
+                    o.write(effected)
+    
+    return wav_fx
+
+
+def generate_and_return_files(bpm, temperature, top_k):
+    midi_events = generate_midi(temperature, top_k)  
+    if not midi_events:
+        return "Error generating MIDI.", None, None
+    
+    write_midi(midi_events, bpm)
+    
+    midi_file = os.path.join(temp_dir, 'output.mid')
+    wav_raw = render_wav(midi_file)
+    wav_fx = os.path.join(temp_dir, 'output_fx.wav')
+
+    sfx_settings = [
+        {
+            'board': Pedalboard([
+                Reverb(room_size=0.01, wet_level=random.uniform(0.005, 0.01), dry_level=0.75, width=1.0),
+                Compressor(threshold_db=-3.0, ratio=8.0, attack_ms=0.0, release_ms=300.0),
+            ])
+        }
+    ]
+
+    for setting in sfx_settings:
+        board = setting['board']
+
+        with AudioFile(wav_raw) as f:
+            with AudioFile(wav_fx, 'w', f.samplerate, f.num_channels) as o:
+                while f.tell() < f.frames:
+                    chunk = f.read(int(f.samplerate))
+                    effected = board(chunk, f.samplerate, reset=False)
+                    o.write(effected)
+
+    return midi_file, wav_fx
+
+
+
+iface = gr.Interface(
+    fn=generate_and_return_files,
+    inputs=[
+        gr.Slider(minimum=60, maximum=200, step=1, value=90, label="bpm"),
+        gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="temperature"),
+        gr.Slider(minimum=4, maximum=128, step=1, value=16, label="top_k")
+    ],
+    outputs=[
+        gr.File(label="MIDI File"),
+        gr.Audio(label="Generated Audio", type="filepath")
+    ],
+    title="<h1 style='font-weight: bold; text-align: center;'>nanoMPC Lite - AI Midi Drum Sequencer</h1>",
+    description="<p style='text-align:center;'>nanoMPC is a tiny transformer model that generates MIDI drum beats inspired by Lo-Fi, Boom Bap and other styles of Hip Hop.</p>"
+)
+
+iface.launch()

checkpoints/model.pt

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

config/eval_gpt2.py

@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=12, n_head=12, n_embd=768
+# 124M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2'

config/eval_gpt2_large.py

@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=36, n_head=20, n_embd=1280
+# 774M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2-large'

config/eval_gpt2_medium.py

@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=24, n_head=16, n_embd=1024
+# 350M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2-medium'

config/eval_gpt2_xl.py

@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=48, n_head=25, n_embd=1600
+# 1558M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2-xl'

config/finetune_shakespeare.py

@@ -0,0 +1,25 @@
+import time
+
+out_dir = 'out-shakespeare'
+eval_interval = 5
+eval_iters = 40
+wandb_log = False # feel free to turn on
+wandb_project = 'shakespeare'
+wandb_run_name = 'ft-' + str(time.time())
+
+dataset = 'shakespeare'
+init_from = 'gpt2-xl' # this is the largest GPT-2 model
+
+# only save checkpoints if the validation loss improves
+always_save_checkpoint = False
+
+# the number of examples per iter:
+# 1 batch_size * 32 grad_accum * 1024 tokens = 32,768 tokens/iter
+# shakespeare has 301,966 tokens, so 1 epoch ~= 9.2 iters
+batch_size = 1
+gradient_accumulation_steps = 32
+max_iters = 20
+
+# finetune at constant LR
+learning_rate = 3e-5
+decay_lr = False

config/train_gpt2.py

@@ -0,0 +1,25 @@
+# config for training GPT-2 (124M) down to very nice loss of ~2.85 on 1 node of 8X A100 40GB
+# launch as the following (e.g. in a screen session) and wait ~5 days:
+# $ torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py
+
+wandb_log = True
+wandb_project = 'owt'
+wandb_run_name='gpt2-124M'
+
+# these make the total batch size be ~0.5M
+# 12 batch size * 1024 block size * 5 gradaccum * 8 GPUs = 491,520
+batch_size = 12
+block_size = 1024
+gradient_accumulation_steps = 5 * 8
+
+# this makes total number of tokens be 300B
+max_iters = 600000
+lr_decay_iters = 600000
+
+# eval stuff
+eval_interval = 1000
+eval_iters = 200
+log_interval = 10
+
+# weight decay
+weight_decay = 1e-1

config/train_shakespeare_char.py

@@ -0,0 +1,37 @@
+# train a miniature character-level shakespeare model
+# good for debugging and playing on macbooks and such
+
+out_dir = 'out-shakespeare-char'
+eval_interval = 250 # keep frequent because we'll overfit
+eval_iters = 200
+log_interval = 10 # don't print too too often
+
+# we expect to overfit on this small dataset, so only save when val improves
+always_save_checkpoint = False
+
+wandb_log = False # override via command line if you like
+wandb_project = 'shakespeare-char'
+wandb_run_name = 'mini-gpt'
+
+dataset = 'shakespeare_char'
+gradient_accumulation_steps = 1
+batch_size = 64
+block_size = 256 # context of up to 256 previous characters
+
+# baby GPT model :)
+n_layer = 6
+n_head = 6
+n_embd = 384
+dropout = 0.2
+
+learning_rate = 1e-3 # with baby networks can afford to go a bit higher
+max_iters = 5000
+lr_decay_iters = 5000 # make equal to max_iters usually
+min_lr = 1e-4 # learning_rate / 10 usually
+beta2 = 0.99 # make a bit bigger because number of tokens per iter is small
+
+warmup_iters = 100 # not super necessary potentially
+
+# on macbook also add
+# device = 'cpu'  # run on cpu only
+# compile = False # do not torch compile the model

configurator.py

@@ -0,0 +1,57 @@
+"""
+Poor Man's Configurator. Probably a terrible idea. Example usage:
+$ python train.py config/override_file.py --batch_size=32
+this will first run config/override_file.py, then override batch_size to 32
+
+The code in this file will be run as follows from e.g. train.py:
+>>> exec(open('configurator.py').read())
+
+So it's not a Python module, it's just shuttling this code away from train.py
+The code in this script then overrides the globals()
+
+I know people are not going to love this, I just really dislike configuration
+complexity and having to prepend config. to every single variable. If someone
+comes up with a better simple Python solution I am all ears.
+"""
+
+import sys
+from ast import literal_eval
+import argparse
+
+parser = argparse.ArgumentParser(description="nanoMPC")
+parser.add_argument("--bpm", type=int, default=90, help="Beats per minute")
+parser.add_argument("--num_samples", type=int, default=1, help="Number of samples")
+args, unknown_args = parser.parse_known_args()  # Capture unknown args for configurator
+
+bpm = args.bpm
+num_samples = args.num_samples
+
+for arg in unknown_args:
+    if arg.startswith('--'):
+        print(f"Skipping command-line argument: {arg}")
+        continue
+
+    if '=' not in arg:
+        config_file = arg
+        print(f"Overriding config with {config_file}:")
+        with open(config_file) as f:
+            print(f.read())
+        exec(open(config_file).read())
+    else:
+        key, val = arg.split('=')
+        key = key[2:]
+
+        if key in globals():
+            if key in ['bpm', 'num_samples']:
+                continue
+            try:
+                attempt = literal_eval(val)
+            except (SyntaxError, ValueError):
+                attempt = val
+            assert type(attempt) == type(globals()[key])
+            print(f"Overriding: {key} = {attempt}")
+            globals()[key] = attempt
+        else:
+            raise ValueError(f"Unknown config key: {key}")
+
+

meta.pkl

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

model.py

@@ -0,0 +1,297 @@
+import inspect
+from dataclasses import dataclass
+from torch.utils.data import DataLoader, Dataset
+from torch.nn import functional as F
+import torch.nn as nn
+import torch
+import math
+import numpy as np
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        self.config = config  # Store the config object
+
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+
+        self.rel_attn_bias = nn.Embedding(config.block_size * 2 - 1, config.n_head)
+
+    def forward(self, x):
+        B, T, C = x.size()
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
+        if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
+            attn_logits = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
+        else:
+            attn_logits = (q @ k.transpose(-2, -1)) / math.sqrt(C // self.n_head)
+            max_rpe = self.config.block_size // 2  # Use config object
+            rpe_matrix = self.generate_rpe(T, max_rpe).to(x.device)
+            rpe_embeddings = self.rel_attn_bias(rpe_matrix).transpose(1, 2).unsqueeze(0)
+            attn_logits = attn_logits + rpe_embeddings
+            attn_logits = attn_logits.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+            attn_logits = F.softmax(attn_logits, dim=-1)
+            attn_logits = self.attn_dropout(attn_logits)
+            attn_logits = attn_logits @ v
+
+        y = attn_logits.transpose(1, 2).contiguous().view(B, T, C)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+    def generate_rpe(self, length, max_rpe):
+        range_vec = torch.arange(length)
+        range_mat = range_vec.unsqueeze(0) - range_vec.unsqueeze(1)
+        range_mat_clipped = torch.clamp(range_mat, -max_rpe, max_rpe)
+        final_mat = range_mat_clipped + max_rpe
+        return final_mat
+
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu    = nn.GELU()
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = RMSNorm(config.n_embd, eps=1e-5)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = RMSNorm(config.n_embd, eps=1e-5)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50304
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True
+
+class GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = RMSNorm(config.n_embd, eps=1e-5),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.transformer.wte.weight = self.lm_head.weight
+
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+        #print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None, noise_pct=0.1):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device)
+
+        tok_emb = self.transformer.wte(idx)
+        pos_emb = self.transformer.wpe(pos)
+
+        # add noise to the input
+        if noise_pct > 0.0:
+            noise_std = torch.std(tok_emb) * noise_pct
+            noise = torch.randn_like(tok_emb) * noise_std
+            tok_emb = tok_emb + noise
+
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+
+            logits = self.lm_head(x[:, [-1], :])
+            loss = None
+
+        return logits, loss
+
+
+    def crop_block_size(self, block_size):
+        assert block_size <= self.config.block_size
+        self.config.block_size = block_size
+        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
+        for block in self.transformer.h:
+            if hasattr(block.attn, 'bias'):
+                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]
+
+    @classmethod
+    def from_pretrained(cls, model_type, override_args=None):
+        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
+        override_args = override_args or {}
+
+        assert all(k == 'dropout' for k in override_args)
+        from transformers import GPT2LMHeadModel
+        print("loading weights from pretrained gpt: %s" % model_type)
+
+        config_args = {
+            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),
+            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024),
+            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280),
+            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600),
+        }[model_type]
+        print("forcing vocab_size=50257, block_size=1024, bias=True")
+        config_args['vocab_size'] = 50257
+        config_args['block_size'] = 1024
+        config_args['bias'] = True
+
+        if 'dropout' in override_args:
+            print(f"overriding dropout rate to {override_args['dropout']}")
+            config_args['dropout'] = override_args['dropout']
+
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = sd.keys()
+        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')]
+
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+
+        sd_keys_hf = sd_hf.keys()
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')]
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')]
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+
+        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        return model
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+
+        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params': decay_params, 'weight_decay': weight_decay},
+            {'params': nodecay_params, 'weight_decay': 0.0}
+        ]
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        extra_args = dict(fused=True) if use_fused else dict()
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
+        print(f"using fused AdamW: {use_fused}")
+
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
+
+        N = self.get_num_params()
+        cfg = self.config
+        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
+        flops_per_token = 6*N + 12*L*H*Q*T
+        flops_per_fwdbwd = flops_per_token * T
+        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
+
+        flops_achieved = flops_per_iter * (1.0/dt)
+        flops_promised = 312e12
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        for _ in range(max_new_tokens):
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx

packages.txt

@@ -0,0 +1 @@
+fluidsynth

requirements.txt

@@ -0,0 +1,4 @@
+pretty_midi==0.2.10
+pedalboard==0.9.3
+torch
+gradio