From 39b1f1c54fe4da7f1d218d11a45363031ac99ffe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Thu, 22 Jun 2023 15:09:03 -0600
Subject: [PATCH 01/84] XX namelist

https://github.com/lfortran/lfortran/issues/1851
---
 driver.f90 | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/driver.f90 b/driver.f90
index 476bd95..0588138 100644
--- a/driver.f90
+++ b/driver.f90
@@ -17,11 +17,11 @@ subroutine load_input(filename, input_txt, n_tokens_to_generate)
 integer, intent(out) :: n_tokens_to_generate
 character(1024) :: input_txt2
 integer :: u, ios
-namelist / input_fastGPT / n_tokens_to_generate
+!namelist / input_fastGPT / n_tokens_to_generate
 allocate(character(0) :: input_txt)
 input_txt = ""
 open(newunit=u, file=filename, status="old")
-read(u, input_fastGPT)
+!read(u, input_fastGPT)
 do
     read(u, "(a)", iostat=ios) input_txt2
     if (ios /= 0) exit

From 5e0df96f4393bbd19086e4bcd48420eb62e060df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Fri, 21 Jul 2023 12:55:29 -0600
Subject: [PATCH 02/84] XX optional array of structures

---
 chat.f90            | 4 +++-
 driver.f90          | 7 ++++---
 tests/test_chat.f90 | 2 +-
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/chat.f90 b/chat.f90
index 9fb6683..b64c7ea 100644
--- a/chat.f90
+++ b/chat.f90
@@ -1,5 +1,7 @@
 program chatgpt2
 use driver, only: chat
+use tokenizer, only: string
 implicit none
-call chat()
+type(string), allocatable :: inputs(:)
+call chat(.false., inputs)
 end program
diff --git a/driver.f90 b/driver.f90
index 0588138..c23e041 100644
--- a/driver.f90
+++ b/driver.f90
@@ -231,7 +231,8 @@ function get_prompt() result(input)
 end if
 end function
 
-subroutine chat(inputs)
+subroutine chat(present_inputs, inputs)
+logical, intent(in) :: present_inputs
 type(string), optional, intent(in) :: inputs(:)
 type(model_t) :: m
 character(:), allocatable :: prompt, input, output
@@ -245,14 +246,14 @@ subroutine chat(inputs)
 &fastGPT: Four." // LF // "&
 &User:"
 write(*,"(a)",advance="no") prompt
-if (present(inputs)) then
+if (present_inputs) then
     n_prompts = size(inputs)
 else
     n_prompts = 1024
 end if
 do i = 1, n_prompts
     write(*,"(a)",advance="no")  " "
-    if (present(inputs)) then
+    if (present_inputs) then
         input = inputs(i)%s
         write(*,"(a)") input
     else
diff --git a/tests/test_chat.f90 b/tests/test_chat.f90
index 41173aa..59fced7 100644
--- a/tests/test_chat.f90
+++ b/tests/test_chat.f90
@@ -19,5 +19,5 @@ program test_chat
     string("What color is snow?"), &
     string("What color do plants usually have?") &
     ]
-call chat(inputs(:3))
+call chat(.true., inputs(:3))
 end program

From 4d1e58f6bc3c83ca6b01a71ea5536fac2129755c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Fri, 21 Jul 2023 18:16:06 -0600
Subject: [PATCH 03/84] XX allocate explicitly

https://github.com/lfortran/lfortran/issues/2022
---
 tokenizer.f90 | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tokenizer.f90 b/tokenizer.f90
index 6bb564c..f669d72 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -123,6 +123,7 @@ function merge_utf8_pairs(intokens) result(tokens)
 type(string), allocatable :: tokens(:)
 integer :: i, j
 logical :: one_more_pass
+allocate(tokens(size(intokens)))
 tokens = intokens
 one_more_pass = .true.
 !print *, "merge_utf8_pairs:", size(tokens)

From f3e0cc8946615764dbf9831ea39dfdd9e43a2cd7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Fri, 21 Jul 2023 18:40:39 -0600
Subject: [PATCH 04/84] XX `print "(1000(i6))", input` not supported

https://github.com/lfortran/lfortran/issues/2024
---
 driver.f90 | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/driver.f90 b/driver.f90
index c23e041..c4fb963 100644
--- a/driver.f90
+++ b/driver.f90
@@ -143,7 +143,8 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print "(a,i4)", "n_tokens_to_generate =", n_tokens_to_generate
 print *
 print "(a)", "Input tokens:"
-print "(1000(i6))", input
+!print "(1000(i6))", input
+print *, input
 print *
 
 if (n_seq + n_tokens_to_generate >= m%n_ctx) then

From 92a8af74e5fe028e5f002c7c27b2aea5f13f9f0c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Fri, 21 Jul 2023 18:57:58 -0600
Subject: [PATCH 05/84] XX `a = [a, b]` array idiom does not work

https://github.com/lfortran/lfortran/issues/2026
---
 tokenizer.f90 | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/tokenizer.f90 b/tokenizer.f90
index f669d72..a89e217 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -186,16 +186,18 @@ function bpe(token, vocab_idx, vocab_txt) result(tokens)
 end function
 
 function encode(input, idx, decoder_txt, vocab_idx, vocab_txt, byte_encoder) &
-        result(tokens)
+        result(tokens2)
 character(*), intent(in) :: input
 integer, intent(in) :: idx(0:), vocab_idx(0:), byte_encoder(0:)
 character, intent(in) :: decoder_txt(:), vocab_txt(:)
-integer, allocatable :: tokens(:)
+integer, parameter :: max_tokens = 2048
+integer :: tokens(max_tokens)
+integer, allocatable :: tokens2(:)
 character(:), allocatable :: tmp, tmp2
 type(string), allocatable :: bpe_tokens(:)
-integer :: i, j, c
+integer :: i, j, c, n_tokens
+n_tokens = 0
 i = 1
-allocate(tokens(0))
 do
     tmp = next_token(input, i)
     if (tmp == "") exit
@@ -210,10 +212,14 @@ function encode(input, idx, decoder_txt, vocab_idx, vocab_txt, byte_encoder) &
     end do
     bpe_tokens = bpe(tmp2, vocab_idx, vocab_txt)
     do j = 1, size(bpe_tokens)
-        tokens = [tokens, word_idx(bpe_tokens(j)%s, idx, decoder_txt)]
+        n_tokens = n_tokens + 1
+        if (n_tokens > max_tokens) error stop "exceeded max_tokens"
+        tokens(n_tokens) = word_idx(bpe_tokens(j)%s, idx, decoder_txt)
     end do
     deallocate(tmp2)
 end do
+allocate(tokens2(n_tokens))
+tokens2 = tokens(:n_tokens)
 end function
 
 function decode(tokens, idx, decoder_txt, byte_decoder) result(output)

From 5fcb41427f4de29cb7d24b4171df0a0b0d7d17cb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Fri, 21 Jul 2023 20:10:26 -0600
Subject: [PATCH 06/84] XX `tokens2 = tokens(:n_tokens)` fails

https://github.com/lfortran/lfortran/issues/2028
---
 tokenizer.f90 | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tokenizer.f90 b/tokenizer.f90
index a89e217..0ad95af 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -219,7 +219,9 @@ function encode(input, idx, decoder_txt, vocab_idx, vocab_txt, byte_encoder) &
     deallocate(tmp2)
 end do
 allocate(tokens2(n_tokens))
-tokens2 = tokens(:n_tokens)
+do i = 1, n_tokens
+    tokens2(i) = tokens(i)
+end do
 end function
 
 function decode(tokens, idx, decoder_txt, byte_decoder) result(output)

From 67df79c5945108506ba2f8e3b930561435d90e58 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Fri, 21 Jul 2023 15:48:33 -0600
Subject: [PATCH 07/84] XX Hardwire the input file

---
 driver.f90 | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/driver.f90 b/driver.f90
index c4fb963..3332cc6 100644
--- a/driver.f90
+++ b/driver.f90
@@ -19,16 +19,17 @@ subroutine load_input(filename, input_txt, n_tokens_to_generate)
 integer :: u, ios
 !namelist / input_fastGPT / n_tokens_to_generate
 allocate(character(0) :: input_txt)
-input_txt = ""
-open(newunit=u, file=filename, status="old")
+n_tokens_to_generate = 20
+input_txt = "Alan Turing theorized that computers would one day become very powerful, but even he could not imagine"
+!open(newunit=u, file=filename, status="old")
 !read(u, input_fastGPT)
-do
-    read(u, "(a)", iostat=ios) input_txt2
-    if (ios /= 0) exit
-    if (len(input_txt) > 0) input_txt = input_txt // char(10)
-    input_txt = input_txt // trim(input_txt2)
-end do
-close(u)
+!do
+!    read(u, "(a)", iostat=ios) input_txt2
+!    if (ios /= 0) exit
+!    if (len(input_txt) > 0) input_txt = input_txt // char(10)
+!    input_txt = input_txt // trim(input_txt2)
+!end do
+!close(u)
 end subroutine
 
 subroutine load_model(filename, m)

From c72c771c326f3de8d4bc3a05d4f47b2c1396eec1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sat, 22 Jul 2023 15:57:35 -0600
Subject: [PATCH 08/84] XX use integer instead of character arrays

This provides a workaround.

https://github.com/lfortran/lfortran/issues/2038
https://github.com/lfortran/lfortran/issues/2041
---
 gpt2.f90      |  2 +-
 tokenizer.f90 | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index a318e5b..443c49e 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -21,7 +21,7 @@ module gpt2_mod
         ln2_b(:,:), ln2_g(:,:), &
         lnf_b(:), lnf_g(:)
     integer, allocatable :: decoder_idx(:), vocab_idx(:), byte_encoder(:)
-    character, allocatable :: decoder_txt(:), vocab_txt(:)
+    integer(1), allocatable :: decoder_txt(:), vocab_txt(:)
     integer :: model_file_version
 end type
 
diff --git a/tokenizer.f90 b/tokenizer.f90
index 0ad95af..1d69736 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -8,12 +8,12 @@ module tokenizer
 contains
 
 function c2s(x) result(y)
-character, intent(in) :: x(:)
+integer(1), intent(in) :: x(:)
 character(:), allocatable :: y
 integer :: i
 allocate(character(size(x)) :: y)
 do i = 1, size(x)
-    y(i:i) = x(i)
+    y(i:i) = char(x(i))
 end do
 end function
 
@@ -62,7 +62,7 @@ function tokenize_word(input, i) result(y)
 function word_idx(word, idx, decoder_txt) result(token)
 character(*), intent(in) :: word
 integer, intent(in) :: idx(0:)
-character, intent(in) :: decoder_txt(:)
+integer(1), intent(in) :: decoder_txt(:)
 integer :: token
 integer :: i
 ! This is O(n) search instead of O(1) lookup in a dictionary, so it is slow
@@ -148,7 +148,7 @@ function bpe(token, vocab_idx, vocab_txt) result(tokens)
 ! Takes a token as a string, and returns bpe tokens as an array of strings
 character(*), intent(in) :: token
 integer, intent(in) :: vocab_idx(0:)
-character, intent(in) :: vocab_txt(:)
+integer(1), intent(in) :: vocab_txt(:)
 type(string), allocatable :: tokens(:)
 integer, allocatable :: pair_scores(:)
 integer :: not_found, merge_pair_idx
@@ -189,7 +189,7 @@ function encode(input, idx, decoder_txt, vocab_idx, vocab_txt, byte_encoder) &
         result(tokens2)
 character(*), intent(in) :: input
 integer, intent(in) :: idx(0:), vocab_idx(0:), byte_encoder(0:)
-character, intent(in) :: decoder_txt(:), vocab_txt(:)
+integer(1), intent(in) :: decoder_txt(:), vocab_txt(:)
 integer, parameter :: max_tokens = 2048
 integer :: tokens(max_tokens)
 integer, allocatable :: tokens2(:)
@@ -226,7 +226,7 @@ function encode(input, idx, decoder_txt, vocab_idx, vocab_txt, byte_encoder) &
 
 function decode(tokens, idx, decoder_txt, byte_decoder) result(output)
 integer, intent(in) :: tokens(:), idx(0:), byte_decoder(0:)
-character, intent(in) :: decoder_txt(:)
+integer(1), intent(in) :: decoder_txt(:)
 character(:), allocatable :: output
 character(:), allocatable :: output2, tmp
 integer :: i, c

From 592c137ac7228bb69abb2fd22b7013342574e6ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sat, 22 Jul 2023 22:24:46 -0600
Subject: [PATCH 09/84] XX char on integer(1)

https://github.com/lfortran/lfortran/issues/2042
---
 tokenizer.f90 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tokenizer.f90 b/tokenizer.f90
index 1d69736..ed7b635 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -13,7 +13,7 @@ function c2s(x) result(y)
 integer :: i
 allocate(character(size(x)) :: y)
 do i = 1, size(x)
-    y(i:i) = char(x(i))
+    y(i:i) = char(int(x(i),4))
 end do
 end function
 

From 97d07df00613b97ab2354d315af5a4fcaef9208d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sat, 22 Jul 2023 23:05:00 -0600
Subject: [PATCH 10/84] XX array indexing bug

https://github.com/lfortran/lfortran/issues/2044
---
 tokenizer.f90 | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tokenizer.f90 b/tokenizer.f90
index ed7b635..d992e9c 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -9,11 +9,13 @@ module tokenizer
 
 function c2s(x) result(y)
 integer(1), intent(in) :: x(:)
+integer(1) :: xx(size(x))
 character(:), allocatable :: y
 integer :: i
+xx = x
 allocate(character(size(x)) :: y)
 do i = 1, size(x)
-    y(i:i) = char(int(x(i),4))
+    y(i:i) = char(int(xx(i),4))
 end do
 end function
 

From cc6888981f11cb57382643ab864c8b90860bf018 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sat, 22 Jul 2023 23:27:02 -0600
Subject: [PATCH 11/84] XX array assignment

https://github.com/lfortran/lfortran/issues/2045
---
 tokenizer.f90 | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tokenizer.f90 b/tokenizer.f90
index d992e9c..928c22c 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -115,8 +115,16 @@ function merge_pair(intokens, idx) result(tokens)
 integer, intent(in) :: idx
 type(string), allocatable :: tokens(:)
 type(string) :: merged_token
+integer :: i
 merged_token%s = intokens(idx)%s // intokens(idx+1)%s
-tokens = [intokens(:idx-1), merged_token, intokens(idx+2:)]
+allocate(tokens(size(intokens)-1))
+do i = 1, idx-1
+    tokens(i) = intokens(i)
+end do
+tokens(idx) = merged_token
+do i = idx+2, size(intokens)
+    tokens(i-1) = intokens(i)
+end do
 end function
 
 function merge_utf8_pairs(intokens) result(tokens)

From e8ced08fbaf790d5d7efc064b0bedc3488060e2e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sat, 22 Jul 2023 23:43:59 -0600
Subject: [PATCH 12/84] XX iachar negative value bug

https://github.com/lfortran/lfortran/issues/2046
---
 tokenizer.f90 | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/tokenizer.f90 b/tokenizer.f90
index 928c22c..775777c 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -131,7 +131,7 @@ function merge_utf8_pairs(intokens) result(tokens)
 ! Merge all UTF-8 character pairs
 type(string), intent(in) :: intokens(:)
 type(string), allocatable :: tokens(:)
-integer :: i, j
+integer :: i, j, ic
 logical :: one_more_pass
 allocate(tokens(size(intokens)))
 tokens = intokens
@@ -142,12 +142,15 @@ function merge_utf8_pairs(intokens) result(tokens)
 do while(one_more_pass)
     one_more_pass = .false.
     do i = j, size(tokens)-1
-        if (len(tokens(i)%s) == 1 .and. iachar(tokens(i)%s(1:1)) >= 128) then
-            tokens = merge_pair(tokens, i)
-            one_more_pass = .true.
-            j = i + 1
-!            print *, "pass"
-            exit
+        if (len(tokens(i)%s) == 1) then
+            ic = iachar(tokens(i)%s(1:1))
+            if (ic < 0) ic = ic + 256
+            if (ic >= 128) then
+                tokens = merge_pair(tokens, i)
+                one_more_pass = .true.
+                j = i + 1
+                exit
+            end if
         end if
     end do
 end do

From bce390971ba6c61c14bed8983fb0db95754fd671 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 00:43:27 -0600
Subject: [PATCH 13/84] XX Remove decoding

---
 driver.f90 | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/driver.f90 b/driver.f90
index 3332cc6..1073489 100644
--- a/driver.f90
+++ b/driver.f90
@@ -154,16 +154,16 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
     error stop
 end if
 
-print "(a)", "Decoded input as text:"
+!print "(a)", "Decoded input as text:"
 !print "(a)", decode(input, decoder_idx, decoder_txt, byte_decoder)
-allocate(character(0) :: output_txt) ! Fix GFortran warning
-output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder)
-print "(a)", output_txt
+!allocate(character(0) :: output_txt) ! Fix GFortran warning
+!output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder)
+!print "(a)", output_txt
 print *
 
-if (input_txt /= output_txt) then
-    error stop "The decoded input text does not agree with the input text"
-end if
+!if (input_txt /= output_txt) then
+!    error stop "The decoded input text does not agree with the input text"
+!end if
 
 allocate(output(n_tokens_to_generate))
 print "(a)", "Running model..."

From d60635ecfec46d8c639c7d83a77d5899fb0b6caf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 00:44:26 -0600
Subject: [PATCH 14/84] XX: array assignment

https://github.com/lfortran/lfortran/issues/2048
---
 gpt2.f90 | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 443c49e..b04eb17 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -210,13 +210,15 @@ function gpt2(n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2,n_layer)
 real(sp) :: y(n_vocab,n_seq_x)
 real(sp) :: x(n_embd,n_seq_x)
-integer :: i
+integer :: i, j
 if (use_kv_cache) then
     i = n_seq
     x(:,1) = wte(:,input(i)+1) + wpe(:,i)
 else
     do i = 1, n_seq
-        x(:,i) = wte(:,input(i)+1) + wpe(:,i)
+        do j = 1, n_embd
+            x(j,i) = wte(j,input(i)+1) + wpe(j,i)
+        end do
     end do
 end if
 do i = 1, n_layer

From 061a61b501e7cdf8f9c0f753aa1485542d099631 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 08:15:51 -0600
Subject: [PATCH 15/84] Add checks for the correct results

---
 driver.f90 | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/driver.f90 b/driver.f90
index 1073489..96e07af 100644
--- a/driver.f90
+++ b/driver.f90
@@ -147,6 +147,11 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 !print "(1000(i6))", input
 print *, input
 print *
+if (n_seq /= 19) error stop
+if (input(1) /= 36235) error stop
+if (input(2) /= 39141) error stop
+if (input(18) /= 407) error stop
+if (input(19) /= 5967) error stop
 
 if (n_seq + n_tokens_to_generate >= m%n_ctx) then
     print *, "The maximum sequence length of the model was surpassed."

From 175e33f32ddfa6fbf38c938bd0fe3f414a5081a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 00:45:04 -0600
Subject: [PATCH 16/84] XX current bug

https://github.com/lfortran/lfortran/issues/2049
---
 gpt2.f90 | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/gpt2.f90 b/gpt2.f90
index b04eb17..e977c0a 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -221,7 +221,9 @@ function gpt2(n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
         end do
     end do
 end if
+print *, "It fails below:"
 do i = 1, n_layer
+    print *, i ! Never gets printed
     x = transformer_block(n_seq, n_seq_x, n_embd, x, &
         mlp_fc_w(:,:,i), mlp_fc_b(:,i), &
         mlp_proj_w(:,:,i), mlp_proj_b(:,i), &

From 42020bced69dec17727f38b13d8fa02bab827291 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 09:48:47 -0600
Subject: [PATCH 17/84] XX comment out transformer

https://github.com/lfortran/lfortran/issues/2049
---
 gpt2.f90 | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index e977c0a..3de7139 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -224,12 +224,12 @@ function gpt2(n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
 print *, "It fails below:"
 do i = 1, n_layer
     print *, i ! Never gets printed
-    x = transformer_block(n_seq, n_seq_x, n_embd, x, &
-        mlp_fc_w(:,:,i), mlp_fc_b(:,i), &
-        mlp_proj_w(:,:,i), mlp_proj_b(:,i), &
-        attn_w(:,:,i), attn_b(:,i), attn_proj_w(:,:,i), attn_proj_b(:,i), &
-        ln1_g(:,i), ln1_b(:,i), ln2_g(:,i), ln2_b(:,i), &
-        n_head, use_kv_cache, kv_cache(:,:,:,i))
+!    x = transformer_block(n_seq, n_seq_x, n_embd, x, &
+!        mlp_fc_w(:,:,i), mlp_fc_b(:,i), &
+!        mlp_proj_w(:,:,i), mlp_proj_b(:,i), &
+!        attn_w(:,:,i), attn_b(:,i), attn_proj_w(:,:,i), attn_proj_b(:,i), &
+!        ln1_g(:,i), ln1_b(:,i), ln2_g(:,i), ln2_b(:,i), &
+!        n_head, use_kv_cache, kv_cache(:,:,:,i))
 end do
 x = layer_norm(x, lnf_g, lnf_b, 1e-5)
 !y = matmul(transpose(wte), x)

From 128baa3f9a459292c1f86852dc4ab1e0b6d44ef1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 09:55:29 -0600
Subject: [PATCH 18/84] XX working on the matmul bug

---
 gpt2.f90          | 1 +
 linalg_c.f90      | 1 +
 linalg_openblas.c | 4 ++++
 3 files changed, 6 insertions(+)

diff --git a/gpt2.f90 b/gpt2.f90
index 3de7139..42bbe28 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -234,6 +234,7 @@ function gpt2(n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
 x = layer_norm(x, lnf_g, lnf_b, 1e-5)
 !y = matmul(transpose(wte), x)
 call matmul_2d_t(wte, x, y)
+stop "OK"
 end function
 
 function generate(n_tokens_to_generate, m, &
diff --git a/linalg_c.f90 b/linalg_c.f90
index 790bb0e..31eb995 100644
--- a/linalg_c.f90
+++ b/linalg_c.f90
@@ -36,6 +36,7 @@ subroutine matmul_2d_t(A, B, C)
     ! C = matmul(transpose(A), B)
     real(sp), intent(in) :: A(:,:), B(:,:)
     real(sp), intent(out) :: C(:,:)
+    print *, "matmul_2d_t"
     call acc_sgemm_t(size(A,2), size(B,2), size(A,1), A, B, C)
     end subroutine
 
diff --git a/linalg_openblas.c b/linalg_openblas.c
index 63637c8..1aad53b 100644
--- a/linalg_openblas.c
+++ b/linalg_openblas.c
@@ -14,5 +14,9 @@ void acc_sgemm_t(int m, int n, int k, float *A, float *B, float *C) {
     //A[k][m] (to be transposed)
     //B[k][n]
     //C[m][n]
+    printf("acc_sgemm_t: %d %d %d\n", m, n, k);
+    // TODO: A seems correct, but B[0] is incorrect:
+    printf("Values: %f %f %f\n", A[0], B[0], C[0]);
     cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, m, n, k, 1.0, A, k, B, k, 0.0, C, m);
+    printf("BLAS done\n");
 }

From da452b5c3c6ef1fabb57b76f972a57fcb4d42ec6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 10:27:55 -0600
Subject: [PATCH 19/84] XX Fix array operations inside a loop

https://github.com/lfortran/lfortran/issues/2053
---
 gpt2.f90 | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 42bbe28..3980ff1 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -63,7 +63,7 @@ function layer_norm(x, g, b, eps) result(y)
 real(sp), intent(in) :: x(:,:), g(:), b(:), eps
 real(sp) :: y(size(x,1),size(x,2))
 real(sp) :: mean(size(x,2)), variance(size(x,2))
-integer :: i
+integer :: i, j
 do i = 1, size(x,2)
     mean(i) = sum(x(:,i)) / size(x,1)
     variance(i) = sum((x(:,i) - mean(i))**2) / size(x,1)
@@ -73,8 +73,10 @@ function layer_norm(x, g, b, eps) result(y)
 !    y(i,:) = g(i) * y(i,:) + b(i)
 !end do
 do i = 1, size(x,2)
-    y(:,i) = (x(:,i) - mean(i)) / sqrt(variance(i) + eps)
-    y(:,i) = g(:) * y(:,i) + b(:)
+    do j = 1, size(x,1)
+        y(j,i) = (x(j,i) - mean(i)) / sqrt(variance(i) + eps)
+        y(j,i) = g(j) * y(j,i) + b(j)
+    end do
 end do
 end function
 

From ad51e32b93afcb31fb1822cf5e045ebf102d0866 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 10:41:32 -0600
Subject: [PATCH 20/84] Move a return array to an argument

https://github.com/lfortran/lfortran/issues/2052
---
 gpt2.f90 | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 3980ff1..b96c990 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -192,12 +192,12 @@ function transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_pr
     mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
 end function
 
-function gpt2(n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
+subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
         wte, wpe, &
         mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, &
         attn_w, attn_b, attn_proj_w, attn_proj_b, &
         ln1_g, ln1_b, ln2_g, ln2_b, lnf_g, lnf_b, &
-        use_kv_cache, kv_cache) result(y)
+        use_kv_cache, kv_cache)
 integer, intent(in) :: n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head
 integer, intent(in) :: input(n_seq)
 real(sp), intent(in) :: wte(n_embd,n_vocab), wpe(n_embd,n_ctx), &
@@ -210,7 +210,7 @@ function gpt2(n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
     lnf_b(n_embd), lnf_g(n_embd)
 logical, intent(in) :: use_kv_cache
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2,n_layer)
-real(sp) :: y(n_vocab,n_seq_x)
+real(sp), intent(out) :: y(n_vocab,n_seq_x)
 real(sp) :: x(n_embd,n_seq_x)
 integer :: i, j
 if (use_kv_cache) then
@@ -237,7 +237,7 @@ function gpt2(n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
 !y = matmul(transpose(wte), x)
 call matmul_2d_t(wte, x, y)
 stop "OK"
-end function
+end subroutine
 
 function generate(n_tokens_to_generate, m, &
         n_seq, input, &
@@ -276,7 +276,7 @@ function generate(n_tokens_to_generate, m, &
         n_seq_x = n_seq2
     end if
     allocate(logits(m%n_vocab, n_seq_x))
-    logits = gpt2(m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, &
+    call gpt2(logits, m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, &
             m%n_head, &
             input2, &
             m%wte, m%wpe, &

From cd297349ea1849c749b62ea1fedc1c788bd6964d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 10:49:32 -0600
Subject: [PATCH 21/84] transformer_block works

---
 gpt2.f90 | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index b96c990..3c4d264 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -173,24 +173,24 @@ function mha(n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head,
 end function
 
 
-function transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, &
+subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, &
         attn_w, attn_b, attn_proj_w, attn_proj_b, ln1_g, ln1_b, ln2_g, ln2_b, &
-        n_head, use_kv_cache, kv_cache) result(y)
-real(sp), intent(in) :: x(n_embd,n_seq_x), &
+        n_head, use_kv_cache, kv_cache)
+real(sp), intent(inout) :: x(n_embd,n_seq_x)
+real(sp), intent(in) :: &
     mlp_fc_w(:,:), mlp_fc_b(:), &
     mlp_proj_w(:,:), mlp_proj_b(:), &
     attn_w(:,:), attn_b(:), attn_proj_w(:,:), attn_proj_b(:), &
     ln1_g(:), ln1_b(:), ln2_g(:), ln2_b(:)
 integer, intent(in) :: n_head
 integer, intent(in) :: n_seq, n_seq_x, n_embd
-real(sp) :: y(n_embd,n_seq_x)
 logical, intent(in) :: use_kv_cache
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
-y = x + mha(n_seq, n_seq_x, n_embd, layer_norm(x, ln1_g, ln1_b, 1e-5_sp), &
+x = x + mha(n_seq, n_seq_x, n_embd, layer_norm(x, ln1_g, ln1_b, 1e-5_sp), &
     attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
-y = y + ffn(layer_norm(y, ln2_g, ln2_b, 1e-5_sp), &
+x = x + ffn(layer_norm(x, ln2_g, ln2_b, 1e-5_sp), &
     mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
-end function
+end subroutine
 
 subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
         wte, wpe, &
@@ -226,12 +226,12 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
 print *, "It fails below:"
 do i = 1, n_layer
     print *, i ! Never gets printed
-!    x = transformer_block(n_seq, n_seq_x, n_embd, x, &
-!        mlp_fc_w(:,:,i), mlp_fc_b(:,i), &
-!        mlp_proj_w(:,:,i), mlp_proj_b(:,i), &
-!        attn_w(:,:,i), attn_b(:,i), attn_proj_w(:,:,i), attn_proj_b(:,i), &
-!        ln1_g(:,i), ln1_b(:,i), ln2_g(:,i), ln2_b(:,i), &
-!        n_head, use_kv_cache, kv_cache(:,:,:,i))
+    call transformer_block(n_seq, n_seq_x, n_embd, x, &
+        mlp_fc_w(:,:,i), mlp_fc_b(:,i), &
+        mlp_proj_w(:,:,i), mlp_proj_b(:,i), &
+        attn_w(:,:,i), attn_b(:,i), attn_proj_w(:,:,i), attn_proj_b(:,i), &
+        ln1_g(:,i), ln1_b(:,i), ln2_g(:,i), ln2_b(:,i), &
+        n_head, use_kv_cache, kv_cache(:,:,:,i))
 end do
 x = layer_norm(x, lnf_g, lnf_b, 1e-5)
 !y = matmul(transpose(wte), x)

From 35c8f7a9e6806d61c54399fead331d3ff7354c9c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 10:51:11 -0600
Subject: [PATCH 22/84] This still works with GF

---
 gpt2.f90          | 6 +++---
 linalg_c.f90      | 2 +-
 linalg_openblas.c | 7 +++----
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 3c4d264..fe4904d 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -223,9 +223,9 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
         end do
     end do
 end if
-print *, "It fails below:"
+!print *, "It fails below:"
 do i = 1, n_layer
-    print *, i ! Never gets printed
+!    print *, i ! Never gets printed
     call transformer_block(n_seq, n_seq_x, n_embd, x, &
         mlp_fc_w(:,:,i), mlp_fc_b(:,i), &
         mlp_proj_w(:,:,i), mlp_proj_b(:,i), &
@@ -236,7 +236,7 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
 x = layer_norm(x, lnf_g, lnf_b, 1e-5)
 !y = matmul(transpose(wte), x)
 call matmul_2d_t(wte, x, y)
-stop "OK"
+!stop "OK"
 end subroutine
 
 function generate(n_tokens_to_generate, m, &
diff --git a/linalg_c.f90 b/linalg_c.f90
index 31eb995..db03267 100644
--- a/linalg_c.f90
+++ b/linalg_c.f90
@@ -36,7 +36,7 @@ subroutine matmul_2d_t(A, B, C)
     ! C = matmul(transpose(A), B)
     real(sp), intent(in) :: A(:,:), B(:,:)
     real(sp), intent(out) :: C(:,:)
-    print *, "matmul_2d_t"
+    !print *, "matmul_2d_t"
     call acc_sgemm_t(size(A,2), size(B,2), size(A,1), A, B, C)
     end subroutine
 
diff --git a/linalg_openblas.c b/linalg_openblas.c
index 1aad53b..2c0c04a 100644
--- a/linalg_openblas.c
+++ b/linalg_openblas.c
@@ -14,9 +14,8 @@ void acc_sgemm_t(int m, int n, int k, float *A, float *B, float *C) {
     //A[k][m] (to be transposed)
     //B[k][n]
     //C[m][n]
-    printf("acc_sgemm_t: %d %d %d\n", m, n, k);
-    // TODO: A seems correct, but B[0] is incorrect:
-    printf("Values: %f %f %f\n", A[0], B[0], C[0]);
+//    printf("acc_sgemm_t: %d %d %d\n", m, n, k);
+//    printf("Values: %f %f %f\n", A[0], B[0], C[0]);
     cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, m, n, k, 1.0, A, k, B, k, 0.0, C, m);
-    printf("BLAS done\n");
+//    printf("BLAS done\n");
 }

From fafe991a31975afb1ee00fa68756678cdeefc98d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:02:17 -0600
Subject: [PATCH 23/84] Rework input handling

---
 gpt2.f90 | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index fe4904d..a587e3f 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -236,7 +236,6 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
 x = layer_norm(x, lnf_g, lnf_b, 1e-5)
 !y = matmul(transpose(wte), x)
 call matmul_2d_t(wte, x, y)
-!stop "OK"
 end subroutine
 
 function generate(n_tokens_to_generate, m, &
@@ -254,22 +253,21 @@ function generate(n_tokens_to_generate, m, &
 integer :: i
 integer :: n_seq2, n_seq_x
 integer :: next_id
-integer, allocatable :: input2(:)
+integer :: input2(size(input)+n_tokens_to_generate)
 logical :: use_kv_cache
 real(sp) :: kv_cache(m%n_embd,n_seq+n_tokens_to_generate,2,m%n_layer)
 character(:), allocatable :: output_txt, last_token
-allocate(input2(size(input)))
 if (present(stop_text)) then
     output_txt = ""
 end if
-input2 = input
+input2(:n_seq) = input
 do i = 1, n_tokens_to_generate
     if (use_cache) then
         use_kv_cache = (i > 1) ! Use cache for subsequent tokens
     else
         use_kv_cache = .false.
     end if
-    n_seq2 = size(input2)
+    n_seq2 = n_seq+i-1
     if (use_kv_cache) then
         n_seq_x = 1
     else
@@ -278,16 +276,16 @@ function generate(n_tokens_to_generate, m, &
     allocate(logits(m%n_vocab, n_seq_x))
     call gpt2(logits, m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, &
             m%n_head, &
-            input2, &
+            input2(:n_seq2), &
             m%wte, m%wpe, &
             m%mlp_fc_w, m%mlp_fc_b, m%mlp_proj_w, m%mlp_proj_b, &
             m%attn_w, m%attn_b, m%attn_proj_w, m%attn_proj_b, &
             m%ln1_g, m%ln1_b, m%ln2_g, m%ln2_b, m%lnf_g, m%lnf_b, use_kv_cache, kv_cache(:,:n_seq2,:,:))
     next_id = maxloc(logits(:,n_seq_x), dim=1)-1
-    input2 = [input2, next_id]
-    last_token = decode([next_id], m%decoder_idx, &
-        m%decoder_txt, byte_decoder)
-    write(*, fmt="(a)", advance="no") last_token
+    input2(n_seq2+1) = next_id
+    !last_token = decode([next_id], m%decoder_idx, &
+    !    m%decoder_txt, byte_decoder)
+    !write(*, fmt="(a)", advance="no") last_token
     if (present(stop_text)) then
         output_txt = output_txt // last_token
         if (output_txt(len(output_txt)-len(stop_text)+1:len(output_txt)) == stop_text) then

From 0c75ecc8575fa5aaac0ce62082e46a5e160c4e52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:05:38 -0600
Subject: [PATCH 24/84] Add a check

---
 driver.f90 | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/driver.f90 b/driver.f90
index 96e07af..ae85fb9 100644
--- a/driver.f90
+++ b/driver.f90
@@ -183,11 +183,15 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print "(a,f8.3,a,f4.2,a)", "    done. Time:", t2o-t1o, "s (", (t2-t1)/(t2o-t1o), "x)"
 print *
 print "(a)", "Output tokens:"
-print "(1000(i6))", output
+print *, output
 output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
 print *
 print "(a)", "Decoded output as text:"
 print "(a)", output_txt
+if (output(1) /= 703) error stop
+if (output(2) /= 484) error stop
+if (output(19) /= 1517) error stop
+if (output(20) /= 318) error stop
 end subroutine
 
 subroutine gpt2_driver3(input_txt, n_tokens_to_generate, stop_text, m, output_txt)

From 1ce230a3703fca11929a686f1b214e2bf2e8745a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:05:59 -0600
Subject: [PATCH 25/84] XX comment out decoding for now

---
 driver.f90 | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/driver.f90 b/driver.f90
index ae85fb9..0168c3e 100644
--- a/driver.f90
+++ b/driver.f90
@@ -184,10 +184,10 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print *
 print "(a)", "Output tokens:"
 print *, output
-output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
-print *
-print "(a)", "Decoded output as text:"
-print "(a)", output_txt
+!output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
+!print *
+!print "(a)", "Decoded output as text:"
+!print "(a)", output_txt
 if (output(1) /= 703) error stop
 if (output(2) /= 484) error stop
 if (output(19) /= 1517) error stop

From 6d52ef3c7b818ffb6f4867975663d1950e8c66a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:11:48 -0600
Subject: [PATCH 26/84] XX use intent out arg

---
 driver.f90 | 4 ++--
 gpt2.f90   | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/driver.f90 b/driver.f90
index 0168c3e..688b5ea 100644
--- a/driver.f90
+++ b/driver.f90
@@ -175,7 +175,7 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 call cpu_time(t1)
 t1o = omp_get_wtime()
 use_cache = .true.
-output = generate(n_tokens_to_generate, m, size(input), input, use_cache, &
+call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
     byte_decoder)
 print *
 t2o = omp_get_wtime()
@@ -225,7 +225,7 @@ subroutine gpt2_driver3(input_txt, n_tokens_to_generate, stop_text, m, output_tx
     error stop "The decoded input text does not agree with the input text"
 end if
 use_cache = .true.
-output = generate(n_tokens_to_generate, m, size(input), input, use_cache, &
+call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
     byte_decoder, stop_text)
 output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
 end subroutine
diff --git a/gpt2.f90 b/gpt2.f90
index a587e3f..bdf455d 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -238,17 +238,17 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
 call matmul_2d_t(wte, x, y)
 end subroutine
 
-function generate(n_tokens_to_generate, m, &
+subroutine generate(output, n_tokens_to_generate, m, &
         n_seq, input, &
         use_cache, &
-        byte_decoder, stop_text) result(output)
+        byte_decoder, stop_text)
 integer, intent(in) :: n_seq, n_tokens_to_generate
 type(model_t), intent(in) :: m
 integer, intent(in) :: input(n_seq)
 logical, intent(in) :: use_cache
 integer, intent(in) :: byte_decoder(:)
 character(*), intent(in), optional :: stop_text ! Stop if you see this text
-integer, allocatable :: output(:)
+integer, allocatable, intent(out) :: output(:)
 real(sp), allocatable :: logits(:,:)
 integer :: i
 integer :: n_seq2, n_seq_x
@@ -295,6 +295,6 @@ function generate(n_tokens_to_generate, m, &
     deallocate(logits)
 end do
 output = input2(n_seq+1:)
-end function
+end subroutine
 
 end module

From 20a245a753cf4f87111f2c62943a3d51829fd0e3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:14:29 -0600
Subject: [PATCH 27/84] XX No allocatable

---
 gpt2.f90 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index bdf455d..fd5377b 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -248,7 +248,7 @@ subroutine generate(output, n_tokens_to_generate, m, &
 logical, intent(in) :: use_cache
 integer, intent(in) :: byte_decoder(:)
 character(*), intent(in), optional :: stop_text ! Stop if you see this text
-integer, allocatable, intent(out) :: output(:)
+integer, intent(out) :: output(:)
 real(sp), allocatable :: logits(:,:)
 integer :: i
 integer :: n_seq2, n_seq_x

From ae12a66451c71130ada06630317457471e511fec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:19:29 -0600
Subject: [PATCH 28/84] XX workaround array op

https://github.com/lfortran/lfortran/issues/2054
---
 gpt2.f90 | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index fd5377b..d7dba20 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -294,7 +294,9 @@ subroutine generate(output, n_tokens_to_generate, m, &
     end if
     deallocate(logits)
 end do
-output = input2(n_seq+1:)
+do i = 1, n_tokens_to_generate
+    output(i) = input2(n_seq+i)
+end do
 end subroutine
 
 end module

From 5d0d1833c54b2662f9d011dc6cb4a67ff8a17fe4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:24:53 -0600
Subject: [PATCH 29/84] XX promote to subroutines

---
 gpt2.f90 | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index d7dba20..a687890 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -92,13 +92,13 @@ function linear(x, w, b) result(y)
 end do
 end function
 
-function ffn(x, fc_w, fc_b, proj_w, proj_b) result(y)
+subroutine ffn(y, x, fc_w, fc_b, proj_w, proj_b)
 real(sp), intent(in) :: x(:,:), fc_w(:,:), fc_b(:), proj_w(:,:), proj_b(:)
-real(sp) :: y(size(x,1),size(x,2))
+real(sp), intent(inout) :: y(size(x,1),size(x,2))
 !real(sp) :: a(4*size(x,1),size(x,2))
 !a = gelu(linear(x, fc_w, fc_b))
-y = linear(gelu(linear(x, fc_w, fc_b)), proj_w, proj_b)
-end function
+y = y + linear(gelu(linear(x, fc_w, fc_b)), proj_w, proj_b)
+end subroutine
 
 function attention(n_embd_head,n_seq,n_seq_x, q, k, v, mask) result(y)
 integer, intent(in) :: n_embd_head, n_seq, n_seq_x
@@ -185,11 +185,12 @@ subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_
 integer, intent(in) :: n_head
 integer, intent(in) :: n_seq, n_seq_x, n_embd
 logical, intent(in) :: use_kv_cache
+real(sp) :: y(n_embd,n_seq_x)
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
 x = x + mha(n_seq, n_seq_x, n_embd, layer_norm(x, ln1_g, ln1_b, 1e-5_sp), &
     attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
-x = x + ffn(layer_norm(x, ln2_g, ln2_b, 1e-5_sp), &
-    mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
+y = layer_norm(x, ln2_g, ln2_b, 1e-5_sp)
+call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
 end subroutine
 
 subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &

From 1ea75852bae47da1d4899f87cb615aab04fa7256 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:37:09 -0600
Subject: [PATCH 30/84] XX layer norm

---
 gpt2.f90 | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index a687890..fba2b5a 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -63,10 +63,14 @@ function layer_norm(x, g, b, eps) result(y)
 real(sp), intent(in) :: x(:,:), g(:), b(:), eps
 real(sp) :: y(size(x,1),size(x,2))
 real(sp) :: mean(size(x,2)), variance(size(x,2))
+real(sp) :: xi(size(x,1))
 integer :: i, j
 do i = 1, size(x,2)
-    mean(i) = sum(x(:,i)) / size(x,1)
-    variance(i) = sum((x(:,i) - mean(i))**2) / size(x,1)
+    do j = 1, size(x,1)
+        xi(j) = x(j,i)
+    end do
+    mean(i) = sum(xi) / size(x,1)
+    variance(i) = sum((xi - mean(i))**2) / size(x,1)
 end do
 !do i = 1, size(x,1)
 !    y(i,:) = (x(i,:) - mean(:)) / sqrt(variance(:) + eps)

From 320f46533227ca4c79aaad9b276d18e56be62c15 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:38:33 -0600
Subject: [PATCH 31/84] XX remove array op

---
 gpt2.f90 | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index fba2b5a..4d65809 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -70,7 +70,10 @@ function layer_norm(x, g, b, eps) result(y)
         xi(j) = x(j,i)
     end do
     mean(i) = sum(xi) / size(x,1)
-    variance(i) = sum((xi - mean(i))**2) / size(x,1)
+    do j = 1, size(x,1)
+        xi(j) = (xi(j) - mean(i))**2
+    end do
+    variance(i) = sum(xi) / size(x,1)
 end do
 !do i = 1, size(x,1)
 !    y(i,:) = (x(i,:) - mean(:)) / sqrt(variance(:) + eps)

From 9ede75687eb1eef720e4b0f1504d4d6785f38a44 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:46:25 -0600
Subject: [PATCH 32/84] XX layer_norm

---
 gpt2.f90 | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 4d65809..f2adeae 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -59,9 +59,9 @@ function softmax(x) result(y)
 end do
 end function
 
-function layer_norm(x, g, b, eps) result(y)
+subroutine layer_norm(y, x, g, b, eps)
 real(sp), intent(in) :: x(:,:), g(:), b(:), eps
-real(sp) :: y(size(x,1),size(x,2))
+real(sp), intent(out) :: y(size(x,1),size(x,2))
 real(sp) :: mean(size(x,2)), variance(size(x,2))
 real(sp) :: xi(size(x,1))
 integer :: i, j
@@ -85,7 +85,7 @@ function layer_norm(x, g, b, eps) result(y)
         y(j,i) = g(j) * y(j,i) + b(j)
     end do
 end do
-end function
+end subroutine
 
 function linear(x, w, b) result(y)
 real(sp), intent(in) :: x(:,:), w(:,:), b(:)
@@ -220,6 +220,7 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2,n_layer)
 real(sp), intent(out) :: y(n_vocab,n_seq_x)
 real(sp) :: x(n_embd,n_seq_x)
+real(sp) :: yy(n_embd,n_seq_x)
 integer :: i, j
 if (use_kv_cache) then
     i = n_seq
@@ -241,7 +242,8 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
         ln1_g(:,i), ln1_b(:,i), ln2_g(:,i), ln2_b(:,i), &
         n_head, use_kv_cache, kv_cache(:,:,:,i))
 end do
-x = layer_norm(x, lnf_g, lnf_b, 1e-5)
+call layer_norm(yy, x, lnf_g, lnf_b, 1e-5)
+x = yy
 !y = matmul(transpose(wte), x)
 call matmul_2d_t(wte, x, y)
 end subroutine

From c5d9c14d53491945d2cf6d7b78f667aa23881dc9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 11:47:00 -0600
Subject: [PATCH 33/84] XX update, works

---
 gpt2.f90 | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index f2adeae..1406cb7 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -194,9 +194,14 @@ subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_
 logical, intent(in) :: use_kv_cache
 real(sp) :: y(n_embd,n_seq_x)
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
-x = x + mha(n_seq, n_seq_x, n_embd, layer_norm(x, ln1_g, ln1_b, 1e-5_sp), &
+call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
+x = x + mha(n_seq, n_seq_x, n_embd, y, &
     attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
-y = layer_norm(x, ln2_g, ln2_b, 1e-5_sp)
+!print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b))
+call layer_norm(y, x, ln2_g, ln2_b, 1e-5_sp)
+!print *, "Out1:", y(1,1)
+!x = y
+!print *, "Out2:", x(1,1)
 call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
 end subroutine
 
@@ -241,7 +246,9 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
         attn_w(:,:,i), attn_b(:,i), attn_proj_w(:,:,i), attn_proj_b(:,i), &
         ln1_g(:,i), ln1_b(:,i), ln2_g(:,i), ln2_b(:,i), &
         n_head, use_kv_cache, kv_cache(:,:,:,i))
+!    print *, x(1,1)
 end do
+!stop "OK"
 call layer_norm(yy, x, lnf_g, lnf_b, 1e-5)
 x = yy
 !y = matmul(transpose(wte), x)

From c409e2a11ee351a5bd100340a84d14528479fbf2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 13:12:13 -0600
Subject: [PATCH 34/84] Works

---
 gpt2.f90 | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 1406cb7..52093ca 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -194,15 +194,18 @@ subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_
 logical, intent(in) :: use_kv_cache
 real(sp) :: y(n_embd,n_seq_x)
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
-call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
-x = x + mha(n_seq, n_seq_x, n_embd, y, &
-    attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
-!print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b))
-call layer_norm(y, x, ln2_g, ln2_b, 1e-5_sp)
-!print *, "Out1:", y(1,1)
-!x = y
-!print *, "Out2:", x(1,1)
-call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
+real(sp) :: tln2_g(size(ln2_g)), tln2_b(size(ln2_b))
+!call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
+!x = x + mha(n_seq, n_seq_x, n_embd, y, &
+!    attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
+tln2_g = ln2_g
+tln2_b = ln2_b
+print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b))
+call layer_norm(y, x, tln2_g, tln2_b, 1e-5_sp)
+print *, "Out1:", y(1,1)
+x = y
+print *, "Out2:", x(1,1)
+!call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
 end subroutine
 
 subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
@@ -248,7 +251,7 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
         n_head, use_kv_cache, kv_cache(:,:,:,i))
 !    print *, x(1,1)
 end do
-!stop "OK"
+stop "OK"
 call layer_norm(yy, x, lnf_g, lnf_b, 1e-5)
 x = yy
 !y = matmul(transpose(wte), x)

From f0342668b43ede5155f084e215f0108470150db5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 13:35:15 -0600
Subject: [PATCH 35/84] Still works

---
 gpt2.f90 | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 52093ca..969c9b3 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -195,17 +195,17 @@ subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_
 real(sp) :: y(n_embd,n_seq_x)
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
 real(sp) :: tln2_g(size(ln2_g)), tln2_b(size(ln2_b))
-!call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
-!x = x + mha(n_seq, n_seq_x, n_embd, y, &
-!    attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
+call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
+x = x + mha(n_seq, n_seq_x, n_embd, y, &
+    attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
 tln2_g = ln2_g
 tln2_b = ln2_b
-print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b))
+!print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b))
 call layer_norm(y, x, tln2_g, tln2_b, 1e-5_sp)
-print *, "Out1:", y(1,1)
-x = y
-print *, "Out2:", x(1,1)
-!call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
+!print *, "Out1:", y(1,1)
+!x = y
+!print *, "Out2:", x(1,1)
+call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
 end subroutine
 
 subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
@@ -251,7 +251,6 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
         n_head, use_kv_cache, kv_cache(:,:,:,i))
 !    print *, x(1,1)
 end do
-stop "OK"
 call layer_norm(yy, x, lnf_g, lnf_b, 1e-5)
 x = yy
 !y = matmul(transpose(wte), x)

From df492af2915e6622ab769c9c9196ec52c4c3e59b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 13:49:48 -0600
Subject: [PATCH 36/84] XX copy arguments

---
 gpt2.f90 | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 969c9b3..826a4c2 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -195,6 +195,10 @@ subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_
 real(sp) :: y(n_embd,n_seq_x)
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
 real(sp) :: tln2_g(size(ln2_g)), tln2_b(size(ln2_b))
+real(sp) :: tmlp_fc_w(size(mlp_fc_w,1),size(mlp_fc_w,2)), &
+    tmlp_fc_b(size(mlp_fc_b)), &
+    tmlp_proj_w(size(mlp_proj_w,1),size(mlp_proj_w,2)), &
+    tmlp_proj_b(size(mlp_proj_b))
 call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
 x = x + mha(n_seq, n_seq_x, n_embd, y, &
     attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
@@ -205,7 +209,11 @@ subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_
 !print *, "Out1:", y(1,1)
 !x = y
 !print *, "Out2:", x(1,1)
-call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
+tmlp_fc_w = mlp_fc_w
+tmlp_fc_b = mlp_fc_b
+tmlp_proj_w = mlp_proj_w
+tmlp_proj_b = mlp_proj_b
+call ffn(x, y, tmlp_fc_w, tmlp_fc_b, tmlp_proj_w, tmlp_proj_b)
 end subroutine
 
 subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &

From 898fc199f8e975d307c4cb5ab1bf7bb7b08ac2ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 14:58:42 -0600
Subject: [PATCH 37/84] Do not do copies

---
 gpt2.f90 | 15 ++-------------
 1 file changed, 2 insertions(+), 13 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 826a4c2..35067f2 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -194,26 +194,15 @@ subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_
 logical, intent(in) :: use_kv_cache
 real(sp) :: y(n_embd,n_seq_x)
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
-real(sp) :: tln2_g(size(ln2_g)), tln2_b(size(ln2_b))
-real(sp) :: tmlp_fc_w(size(mlp_fc_w,1),size(mlp_fc_w,2)), &
-    tmlp_fc_b(size(mlp_fc_b)), &
-    tmlp_proj_w(size(mlp_proj_w,1),size(mlp_proj_w,2)), &
-    tmlp_proj_b(size(mlp_proj_b))
 call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
 x = x + mha(n_seq, n_seq_x, n_embd, y, &
     attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
-tln2_g = ln2_g
-tln2_b = ln2_b
 !print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b))
-call layer_norm(y, x, tln2_g, tln2_b, 1e-5_sp)
+call layer_norm(y, x, ln2_g, ln2_b, 1e-5_sp)
 !print *, "Out1:", y(1,1)
 !x = y
 !print *, "Out2:", x(1,1)
-tmlp_fc_w = mlp_fc_w
-tmlp_fc_b = mlp_fc_b
-tmlp_proj_w = mlp_proj_w
-tmlp_proj_b = mlp_proj_b
-call ffn(x, y, tmlp_fc_w, tmlp_fc_b, tmlp_proj_w, tmlp_proj_b)
+call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
 end subroutine
 
 subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &

From 33dc91baee700c1cec30503fb146e731585c8080 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 14:59:26 -0600
Subject: [PATCH 38/84] Works

---
 gpt2.f90 | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 35067f2..60ae96d 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -238,7 +238,9 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
     end do
 end if
 !print *, "It fails below:"
-do i = 1, n_layer
+do j = 1, n_layer
+    i = j
+    !i = 1
 !    print *, i ! Never gets printed
     call transformer_block(n_seq, n_seq_x, n_embd, x, &
         mlp_fc_w(:,:,i), mlp_fc_b(:,i), &

From c8b2b4fd0162f5361fe735b8fd56b16bae3014e8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:09:01 -0600
Subject: [PATCH 39/84] Simplify ffn

---
 gpt2.f90 | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 60ae96d..4c10c63 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -90,21 +90,27 @@ subroutine layer_norm(y, x, g, b, eps)
 function linear(x, w, b) result(y)
 real(sp), intent(in) :: x(:,:), w(:,:), b(:)
 real(sp) :: y(size(b,1),size(x,2))
-integer :: i
+integer :: i, j
 !y = matmul(w, x) + spread(b, 2, size(x,2))
 !y = matmul(w, x)
 call matmul_2d(w, x, y)
 do i = 1, size(y,2)
-    y(:,i) = y(:,i) + b(:)
+    do j = 1, size(y,1)
+        y(j,i) = y(j,i) + b(j)
+    end do
 end do
 end function
 
 subroutine ffn(y, x, fc_w, fc_b, proj_w, proj_b)
 real(sp), intent(in) :: x(:,:), fc_w(:,:), fc_b(:), proj_w(:,:), proj_b(:)
 real(sp), intent(inout) :: y(size(x,1),size(x,2))
-!real(sp) :: a(4*size(x,1),size(x,2))
-!a = gelu(linear(x, fc_w, fc_b))
-y = y + linear(gelu(linear(x, fc_w, fc_b)), proj_w, proj_b)
+real(sp) :: yy(size(x,1),size(x,2))
+real(sp) :: a(4*size(x,1),size(x,2))
+real(sp) :: aa(4*size(x,1),size(x,2))
+aa = linear(x, fc_w, fc_b)
+a = gelu(aa)
+yy = linear(a, proj_w, proj_b)
+y = y + yy
 end subroutine
 
 function attention(n_embd_head,n_seq,n_seq_x, q, k, v, mask) result(y)

From 093eae03936c670b40516cd43cf5ba60f3f2fadc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:13:25 -0600
Subject: [PATCH 40/84] Rework mha

---
 gpt2.f90 | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 4c10c63..39ef1f1 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -124,9 +124,8 @@ function attention(n_embd_head,n_seq,n_seq_x, q, k, v, mask) result(y)
 call matmul_2d(v, softmax(tmp / sqrt(real(n_embd_head,sp)) + mask), y)
 end function
 
-function mha(n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head, &
-            use_kv_cache, kv_cache) &
-        result(y)
+subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head, &
+            use_kv_cache, kv_cache)
 integer, intent(in) :: n_seq, n_seq_x, n_embd
 real(sp), intent(in) :: x(n_embd,n_seq_x), &
     attn_w(3*n_embd,n_embd), attn_b(3*n_embd), &
@@ -134,7 +133,7 @@ function mha(n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head,
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
 integer, intent(in) :: n_head
 logical, intent(in) :: use_kv_cache
-real(sp) :: y(n_embd,n_seq_x)
+real(sp), intent(out) :: y(n_embd,n_seq_x)
 real(sp) :: causal_mask(n_seq,n_seq_x)
 real(sp) :: x2(3*n_embd,n_seq_x)
 integer :: i, j
@@ -183,7 +182,7 @@ function mha(n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head,
 end associate
 ! Out projection
 y = linear(y, proj_w, proj_b)
-end function
+end subroutine
 
 
 subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, &
@@ -199,10 +198,12 @@ subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_
 integer, intent(in) :: n_seq, n_seq_x, n_embd
 logical, intent(in) :: use_kv_cache
 real(sp) :: y(n_embd,n_seq_x)
+real(sp) :: yy(n_embd,n_seq_x)
 real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
 call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
-x = x + mha(n_seq, n_seq_x, n_embd, y, &
+call mha(yy, n_seq, n_seq_x, n_embd, y, &
     attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
+x = x + yy
 !print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b))
 call layer_norm(y, x, ln2_g, ln2_b, 1e-5_sp)
 !print *, "Out1:", y(1,1)

From cd4ed66a3da6a0282a146fdae656b6dd729ae297 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:23:52 -0600
Subject: [PATCH 41/84] Disable kv_cache for now

---
 driver.f90 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/driver.f90 b/driver.f90
index 688b5ea..c2a7d15 100644
--- a/driver.f90
+++ b/driver.f90
@@ -174,7 +174,7 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print "(a)", "Running model..."
 call cpu_time(t1)
 t1o = omp_get_wtime()
-use_cache = .true.
+use_cache = .false.
 call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
     byte_decoder)
 print *

From a4b2975454aa43feea0d7a7bbc8ad6cd963769c1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:27:23 -0600
Subject: [PATCH 42/84] Get rid of associate

---
 gpt2.f90 | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 39ef1f1..0ada45a 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -152,19 +152,18 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
     end do
 end if
 x2 = linear(x, attn_w, attn_b)
-associate ( &
-        q => x2((1-1)*n_embd+1:1*n_embd,:), &
-        k => x2((2-1)*n_embd+1:2*n_embd,:), &
-        v => x2((3-1)*n_embd+1:3*n_embd,:)  &
-    )
-    if (use_kv_cache) then
-        kv_cache(:,n_seq,1) = k(:,1)
-        kv_cache(:,n_seq,2) = v(:,1)
-    else
-        kv_cache(:,:,1) = k
-        kv_cache(:,:,2) = v
-    end if
-end associate
+if (use_kv_cache) then
+    error stop
+    !kv_cache(:,n_seq,1) = k(:,1)
+    !kv_cache(:,n_seq,2) = v(:,1)
+else
+    do i = 1, n_seq
+    do j = 1, n_embd
+        kv_cache(j,i,1) = x2((2-1)*n_embd+j,i)
+        kv_cache(j,i,2) = x2((3-1)*n_embd+j,i)
+    end do
+    end do
+end if
 associate ( &
         q => x2((1-1)*n_embd+1:1*n_embd,:), &
         k => kv_cache(:,:,1), &

From 5fe8374057d5cce5ea0ac5105ca61515325202b2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:28:57 -0600
Subject: [PATCH 43/84] Get kv cache working again

---
 driver.f90 | 2 +-
 gpt2.f90   | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/driver.f90 b/driver.f90
index c2a7d15..688b5ea 100644
--- a/driver.f90
+++ b/driver.f90
@@ -174,7 +174,7 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print "(a)", "Running model..."
 call cpu_time(t1)
 t1o = omp_get_wtime()
-use_cache = .false.
+use_cache = .true.
 call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
     byte_decoder)
 print *
diff --git a/gpt2.f90 b/gpt2.f90
index 0ada45a..750f16f 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -153,9 +153,10 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
 end if
 x2 = linear(x, attn_w, attn_b)
 if (use_kv_cache) then
-    error stop
-    !kv_cache(:,n_seq,1) = k(:,1)
-    !kv_cache(:,n_seq,2) = v(:,1)
+    do j = 1, n_embd
+        kv_cache(j,n_seq,1) = x2((2-1)*n_embd+j,1)
+        kv_cache(j,n_seq,2) = x2((3-1)*n_embd+j,1)
+    end do
 else
     do i = 1, n_seq
     do j = 1, n_embd

From 0da395486daa1c096947eb4f185c1dce17e43fa2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:35:10 -0600
Subject: [PATCH 44/84] Simplify

---
 gpt2.f90 | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 750f16f..005ff48 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -165,18 +165,14 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
     end do
     end do
 end if
-associate ( &
-        q => x2((1-1)*n_embd+1:1*n_embd,:), &
-        k => kv_cache(:,:,1), &
-        v => kv_cache(:,:,2)  &
-    )
+associate ( q => x2((1-1)*n_embd+1:1*n_embd,:) )
     ! Perform attention over each head
     do i = 1, n_head
         y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:) = attention( &
             n_embd/n_head, n_seq, n_seq_x, &
             q((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
-            k((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
-            v((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
+            kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,1), &
+            kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,2), &
             causal_mask)
     end do
 end associate

From abb8d10967f96e1e0d483723100d71675dc5c604 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:37:04 -0600
Subject: [PATCH 45/84] Get rid of the second associate

---
 gpt2.f90 | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 005ff48..a7ba4d5 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -165,17 +165,15 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
     end do
     end do
 end if
-associate ( q => x2((1-1)*n_embd+1:1*n_embd,:) )
-    ! Perform attention over each head
-    do i = 1, n_head
-        y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:) = attention( &
-            n_embd/n_head, n_seq, n_seq_x, &
-            q((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
-            kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,1), &
-            kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,2), &
-            causal_mask)
-    end do
-end associate
+! Perform attention over each head
+do i = 1, n_head
+    y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:) = attention( &
+        n_embd/n_head, n_seq, n_seq_x, &
+        x2((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
+        kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,1), &
+        kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,2), &
+        causal_mask)
+end do
 ! Out projection
 y = linear(y, proj_w, proj_b)
 end subroutine

From 0f1461cb1f3031cf13059b1bcea55b2d4242380e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:38:42 -0600
Subject: [PATCH 46/84] Make attention a subroutine

---
 gpt2.f90 | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index a7ba4d5..46ff30b 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -113,16 +113,16 @@ subroutine ffn(y, x, fc_w, fc_b, proj_w, proj_b)
 y = y + yy
 end subroutine
 
-function attention(n_embd_head,n_seq,n_seq_x, q, k, v, mask) result(y)
+subroutine attention(y, n_embd_head,n_seq,n_seq_x, q, k, v, mask)
 integer, intent(in) :: n_embd_head, n_seq, n_seq_x
 real(sp), intent(in) :: q(n_embd_head,n_seq_x), k(n_embd_head,n_seq), v(n_embd_head,n_seq), mask(n_seq,n_seq_x)
-real(sp) :: y(n_embd_head,n_seq_x)
+real(sp), intent(out) :: y(n_embd_head,n_seq_x)
 real(sp) :: tmp(n_seq,n_seq_x)
 !tmp = matmul(transpose(k), q)
 !call matmul_2d(transpose(k), q, tmp)
 call matmul_2d_t(k, q, tmp)
 call matmul_2d(v, softmax(tmp / sqrt(real(n_embd_head,sp)) + mask), y)
-end function
+end subroutine
 
 subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head, &
             use_kv_cache, kv_cache)
@@ -167,7 +167,7 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
 end if
 ! Perform attention over each head
 do i = 1, n_head
-    y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:) = attention( &
+    call attention(y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
         n_embd/n_head, n_seq, n_seq_x, &
         x2((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
         kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,1), &

From 4a1527ef0dec214e6a8d79b72c096c9c7b5bac17 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:45:30 -0600
Subject: [PATCH 47/84] Use local variables

---
 gpt2.f90 | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 46ff30b..b4313b3 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -136,6 +136,7 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
 real(sp), intent(out) :: y(n_embd,n_seq_x)
 real(sp) :: causal_mask(n_seq,n_seq_x)
 real(sp) :: x2(3*n_embd,n_seq_x)
+real(sp) :: q(n_embd/n_head,n_seq_x), k(n_embd/n_head,n_seq), v(n_embd/n_head,n_seq)
 integer :: i, j
 ! Mask
 if (use_kv_cache) then
@@ -167,12 +168,11 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
 end if
 ! Perform attention over each head
 do i = 1, n_head
+    q = x2((i-1)*n_embd/n_head+1:i*n_embd/n_head,:)
+    k = kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,1)
+    v = kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,2)
     call attention(y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
-        n_embd/n_head, n_seq, n_seq_x, &
-        x2((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
-        kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,1), &
-        kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,2), &
-        causal_mask)
+        n_embd/n_head, n_seq, n_seq_x, q, k, v, causal_mask)
 end do
 ! Out projection
 y = linear(y, proj_w, proj_b)

From 4c88fe70c6f9518b7f355ad74440ead386a7c98c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:52:37 -0600
Subject: [PATCH 48/84] y split

---
 gpt2.f90 | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index b4313b3..0e00089 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -137,6 +137,7 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
 real(sp) :: causal_mask(n_seq,n_seq_x)
 real(sp) :: x2(3*n_embd,n_seq_x)
 real(sp) :: q(n_embd/n_head,n_seq_x), k(n_embd/n_head,n_seq), v(n_embd/n_head,n_seq)
+real(sp) :: yy(n_embd/n_head,n_seq_x)
 integer :: i, j
 ! Mask
 if (use_kv_cache) then
@@ -171,8 +172,8 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
     q = x2((i-1)*n_embd/n_head+1:i*n_embd/n_head,:)
     k = kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,1)
     v = kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,2)
-    call attention(y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), &
-        n_embd/n_head, n_seq, n_seq_x, q, k, v, causal_mask)
+    call attention(yy, n_embd/n_head, n_seq, n_seq_x, q, k, v, causal_mask)
+    y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:) = yy
 end do
 ! Out projection
 y = linear(y, proj_w, proj_b)

From d930f0b87e8788bd009b0024c8f741b4152ae132 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:54:47 -0600
Subject: [PATCH 49/84] Use l

---
 gpt2.f90 | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 0e00089..b5e0077 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -138,7 +138,7 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
 real(sp) :: x2(3*n_embd,n_seq_x)
 real(sp) :: q(n_embd/n_head,n_seq_x), k(n_embd/n_head,n_seq), v(n_embd/n_head,n_seq)
 real(sp) :: yy(n_embd/n_head,n_seq_x)
-integer :: i, j
+integer :: i, j, l
 ! Mask
 if (use_kv_cache) then
     causal_mask = 0
@@ -168,12 +168,12 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
     end do
 end if
 ! Perform attention over each head
-do i = 1, n_head
-    q = x2((i-1)*n_embd/n_head+1:i*n_embd/n_head,:)
-    k = kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,1)
-    v = kv_cache((i-1)*n_embd/n_head+1:i*n_embd/n_head,:,2)
+do l = 1, n_head
+    q = x2((l-1)*n_embd/n_head+1:l*n_embd/n_head,:)
+    k = kv_cache((l-1)*n_embd/n_head+1:l*n_embd/n_head,:,1)
+    v = kv_cache((l-1)*n_embd/n_head+1:l*n_embd/n_head,:,2)
     call attention(yy, n_embd/n_head, n_seq, n_seq_x, q, k, v, causal_mask)
-    y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:) = yy
+    y((l-1)*n_embd/n_head+1:l*n_embd/n_head,:) = yy
 end do
 ! Out projection
 y = linear(y, proj_w, proj_b)

From d8678570424266b3a15fa39fe4fafea6171af0ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:56:26 -0600
Subject: [PATCH 50/84] Add a loop

---
 gpt2.f90 | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index b5e0077..2a4098b 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -170,8 +170,12 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
 ! Perform attention over each head
 do l = 1, n_head
     q = x2((l-1)*n_embd/n_head+1:l*n_embd/n_head,:)
-    k = kv_cache((l-1)*n_embd/n_head+1:l*n_embd/n_head,:,1)
-    v = kv_cache((l-1)*n_embd/n_head+1:l*n_embd/n_head,:,2)
+    do i = 1, n_seq
+    do j = 1, n_embd/n_head
+        k(j,i) = kv_cache((l-1)*n_embd/n_head+j,i,1)
+        v(j,i) = kv_cache((l-1)*n_embd/n_head+j,i,2)
+    end do
+    end do
     call attention(yy, n_embd/n_head, n_seq, n_seq_x, q, k, v, causal_mask)
     y((l-1)*n_embd/n_head+1:l*n_embd/n_head,:) = yy
 end do

From 77bcc186aa0f16a621c187c23f57cd9b6657c1f4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:57:34 -0600
Subject: [PATCH 51/84] Loop

---
 gpt2.f90 | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 2a4098b..34e46a0 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -169,7 +169,11 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
 end if
 ! Perform attention over each head
 do l = 1, n_head
-    q = x2((l-1)*n_embd/n_head+1:l*n_embd/n_head,:)
+    do i = 1, n_seq_x
+    do j = 1, n_embd/n_head
+        q(j,i) = x2((l-1)*n_embd/n_head+j,i)
+    end do
+    end do
     do i = 1, n_seq
     do j = 1, n_embd/n_head
         k(j,i) = kv_cache((l-1)*n_embd/n_head+j,i,1)

From ea3f07cfb210ff908abd1c6772090bbc8160fafe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 15:58:37 -0600
Subject: [PATCH 52/84] Loop

---
 gpt2.f90 | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 34e46a0..6a4c0d0 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -181,7 +181,11 @@ subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_h
     end do
     end do
     call attention(yy, n_embd/n_head, n_seq, n_seq_x, q, k, v, causal_mask)
-    y((l-1)*n_embd/n_head+1:l*n_embd/n_head,:) = yy
+    do i = 1, n_seq_x
+    do j = 1, n_embd/n_head
+        y((l-1)*n_embd/n_head+j,i) = yy(j,i)
+    end do
+    end do
 end do
 ! Out projection
 y = linear(y, proj_w, proj_b)

From c4524b2b00678c075885f99bd67ddd6892b79a03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 16:01:37 -0600
Subject: [PATCH 53/84] Separate

---
 gpt2.f90 | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 6a4c0d0..6b63d41 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -121,7 +121,9 @@ subroutine attention(y, n_embd_head,n_seq,n_seq_x, q, k, v, mask)
 !tmp = matmul(transpose(k), q)
 !call matmul_2d(transpose(k), q, tmp)
 call matmul_2d_t(k, q, tmp)
-call matmul_2d(v, softmax(tmp / sqrt(real(n_embd_head,sp)) + mask), y)
+tmp = tmp / sqrt(real(n_embd_head,sp)) + mask
+tmp = softmax(tmp)
+call matmul_2d(v, tmp, y)
 end subroutine
 
 subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head, &

From 4a6876cd06139dd8f26b2746ff53c2b495930975 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 16:02:36 -0600
Subject: [PATCH 54/84] Loop

---
 gpt2.f90 | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 6b63d41..43d133f 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -118,10 +118,15 @@ subroutine attention(y, n_embd_head,n_seq,n_seq_x, q, k, v, mask)
 real(sp), intent(in) :: q(n_embd_head,n_seq_x), k(n_embd_head,n_seq), v(n_embd_head,n_seq), mask(n_seq,n_seq_x)
 real(sp), intent(out) :: y(n_embd_head,n_seq_x)
 real(sp) :: tmp(n_seq,n_seq_x)
+integer :: i, j
 !tmp = matmul(transpose(k), q)
 !call matmul_2d(transpose(k), q, tmp)
 call matmul_2d_t(k, q, tmp)
-tmp = tmp / sqrt(real(n_embd_head,sp)) + mask
+do i = 1, n_seq_x
+do j = 1, n_seq
+    tmp(j,i) = tmp(j,i) / sqrt(real(n_embd_head,sp)) + mask(j,i)
+end do
+end do
 tmp = softmax(tmp)
 call matmul_2d(v, tmp, y)
 end subroutine

From 6de7e980893fd00b80a59acb95f406d4693dcb75 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 16:05:25 -0600
Subject: [PATCH 55/84] Loops

---
 gpt2.f90 | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 43d133f..1bcd23f 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -52,10 +52,16 @@ elemental real(sp) function gelu(x) result(y)
 function softmax(x) result(y)
 real(sp), intent(in) :: x(:,:)
 real(sp) :: y(size(x,1),size(x,2))
-integer :: i
+integer :: i,j
+real(sp) :: s
 do i = 1, size(x,2)
-    y(:,i) = exp(x(:,i) - maxval(x(:,i)))
-    y(:,i) = y(:,i) / sum(y(:,i))
+    do j = 1, size(x,1)
+        y(j,i) = exp(x(j,i) - maxval(x(:,i)))
+    end do
+    s = sum(y(:,i))
+    do j = 1, size(x,1)
+        y(j,i) = y(j,i) / s
+    end do
 end do
 end function
 

From 6c5591c5813305d86a0a352d6226e22acf2d4d48 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 16:06:12 -0600
Subject: [PATCH 56/84] Loop

---
 gpt2.f90 | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 1bcd23f..70367d2 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -58,7 +58,10 @@ function softmax(x) result(y)
     do j = 1, size(x,1)
         y(j,i) = exp(x(j,i) - maxval(x(:,i)))
     end do
-    s = sum(y(:,i))
+    s = 0
+    do j = 1, size(x,1)
+        s = s + y(j,i)
+    end do
     do j = 1, size(x,1)
         y(j,i) = y(j,i) / s
     end do

From a5e35e7541bb5ffcffd4357f84b73fe78023eda4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 16:07:55 -0600
Subject: [PATCH 57/84] Loops

---
 gpt2.f90 | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 70367d2..e04fec5 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -55,8 +55,12 @@ function softmax(x) result(y)
 integer :: i,j
 real(sp) :: s
 do i = 1, size(x,2)
+    s = -1e10
     do j = 1, size(x,1)
-        y(j,i) = exp(x(j,i) - maxval(x(:,i)))
+        if (x(j,i) > s) s = x(j,i)
+    end do
+    do j = 1, size(x,1)
+        y(j,i) = exp(x(j,i) - s)
     end do
     s = 0
     do j = 1, size(x,1)

From c2f66296136022cf2fa622725fc14febeafbbf0a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Sun, 23 Jul 2023 23:55:00 -0600
Subject: [PATCH 58/84] Disable cache

---
 driver.f90 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/driver.f90 b/driver.f90
index 688b5ea..c2a7d15 100644
--- a/driver.f90
+++ b/driver.f90
@@ -174,7 +174,7 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print "(a)", "Running model..."
 call cpu_time(t1)
 t1o = omp_get_wtime()
-use_cache = .true.
+use_cache = .false.
 call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
     byte_decoder)
 print *

From 33ae02661f10762a4bb101a18ac625911c32fe79 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 00:01:56 -0600
Subject: [PATCH 59/84] Enable cache

---
 driver.f90 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/driver.f90 b/driver.f90
index c2a7d15..688b5ea 100644
--- a/driver.f90
+++ b/driver.f90
@@ -174,7 +174,7 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print "(a)", "Running model..."
 call cpu_time(t1)
 t1o = omp_get_wtime()
-use_cache = .false.
+use_cache = .true.
 call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
     byte_decoder)
 print *

From bcac8b2e50a45ba48529523eee44480c06a6a972 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 00:10:10 -0600
Subject: [PATCH 60/84] Temp copy

---
 gpt2.f90 | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index e04fec5..ac72a9e 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -263,7 +263,9 @@ subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, inpu
 integer :: i, j
 if (use_kv_cache) then
     i = n_seq
-    x(:,1) = wte(:,input(i)+1) + wpe(:,i)
+    do j = 1, n_embd
+        x(j,1) = wte(j,input(i)+1) + wpe(j,i)
+    end do
 else
     do i = 1, n_seq
         do j = 1, n_embd
@@ -308,6 +310,7 @@ subroutine generate(output, n_tokens_to_generate, m, &
 integer :: input2(size(input)+n_tokens_to_generate)
 logical :: use_kv_cache
 real(sp) :: kv_cache(m%n_embd,n_seq+n_tokens_to_generate,2,m%n_layer)
+real(sp), allocatable :: kv_cache2(:,:,:,:)
 character(:), allocatable :: output_txt, last_token
 if (present(stop_text)) then
     output_txt = ""
@@ -326,13 +329,18 @@ subroutine generate(output, n_tokens_to_generate, m, &
         n_seq_x = n_seq2
     end if
     allocate(logits(m%n_vocab, n_seq_x))
+    allocate(kv_cache2(m%n_embd,n_seq+n_tokens_to_generate,2,m%n_layer))
+    kv_cache2 = kv_cache(:,:n_seq2,:,:)
     call gpt2(logits, m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, &
             m%n_head, &
             input2(:n_seq2), &
             m%wte, m%wpe, &
             m%mlp_fc_w, m%mlp_fc_b, m%mlp_proj_w, m%mlp_proj_b, &
             m%attn_w, m%attn_b, m%attn_proj_w, m%attn_proj_b, &
-            m%ln1_g, m%ln1_b, m%ln2_g, m%ln2_b, m%lnf_g, m%lnf_b, use_kv_cache, kv_cache(:,:n_seq2,:,:))
+            m%ln1_g, m%ln1_b, m%ln2_g, m%ln2_b, m%lnf_g, m%lnf_b, use_kv_cache,&
+            kv_cache2)
+    kv_cache(:,:n_seq2,:,:) = kv_cache2
+    deallocate(kv_cache2)
     next_id = maxloc(logits(:,n_seq_x), dim=1)-1
     input2(n_seq2+1) = next_id
     !last_token = decode([next_id], m%decoder_idx, &

From 13f802130f43b94b1a347960f6b8034914cc4606 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 00:11:14 -0600
Subject: [PATCH 61/84] Align ranks

---
 gpt2.f90 | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index ac72a9e..668d05d 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -329,8 +329,8 @@ subroutine generate(output, n_tokens_to_generate, m, &
         n_seq_x = n_seq2
     end if
     allocate(logits(m%n_vocab, n_seq_x))
-    allocate(kv_cache2(m%n_embd,n_seq+n_tokens_to_generate,2,m%n_layer))
-    kv_cache2 = kv_cache(:,:n_seq2,:,:)
+    allocate(kv_cache2(m%n_embd,n_seq2,2,m%n_layer))
+    kv_cache2(:,:,:,:) = kv_cache(:,:n_seq2,:,:)
     call gpt2(logits, m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, &
             m%n_head, &
             input2(:n_seq2), &
@@ -339,7 +339,7 @@ subroutine generate(output, n_tokens_to_generate, m, &
             m%attn_w, m%attn_b, m%attn_proj_w, m%attn_proj_b, &
             m%ln1_g, m%ln1_b, m%ln2_g, m%ln2_b, m%lnf_g, m%lnf_b, use_kv_cache,&
             kv_cache2)
-    kv_cache(:,:n_seq2,:,:) = kv_cache2
+    kv_cache(:,:n_seq2,:,:) = kv_cache2(:,:,:,:)
     deallocate(kv_cache2)
     next_id = maxloc(logits(:,n_seq_x), dim=1)-1
     input2(n_seq2+1) = next_id

From 23efcd60f404e9e0a0c3da4e7e120a10d653eb93 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 00:13:52 -0600
Subject: [PATCH 62/84] Copy things element by element

---
 gpt2.f90 | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index 668d05d..e1db77d 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -304,7 +304,7 @@ subroutine generate(output, n_tokens_to_generate, m, &
 character(*), intent(in), optional :: stop_text ! Stop if you see this text
 integer, intent(out) :: output(:)
 real(sp), allocatable :: logits(:,:)
-integer :: i
+integer :: i, i1, i2, i3, i4
 integer :: n_seq2, n_seq_x
 integer :: next_id
 integer :: input2(size(input)+n_tokens_to_generate)
@@ -330,7 +330,15 @@ subroutine generate(output, n_tokens_to_generate, m, &
     end if
     allocate(logits(m%n_vocab, n_seq_x))
     allocate(kv_cache2(m%n_embd,n_seq2,2,m%n_layer))
-    kv_cache2(:,:,:,:) = kv_cache(:,:n_seq2,:,:)
+    do i4 = 1, m%n_layer
+    do i3 = 1, 2
+    do i2 = 1, n_seq2
+    do i1 = 1, m%n_embd
+        kv_cache2(i1,i2,i3,i4) = kv_cache(i1,i2,i3,i4)
+    end do
+    end do
+    end do
+    end do
     call gpt2(logits, m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, &
             m%n_head, &
             input2(:n_seq2), &
@@ -339,7 +347,15 @@ subroutine generate(output, n_tokens_to_generate, m, &
             m%attn_w, m%attn_b, m%attn_proj_w, m%attn_proj_b, &
             m%ln1_g, m%ln1_b, m%ln2_g, m%ln2_b, m%lnf_g, m%lnf_b, use_kv_cache,&
             kv_cache2)
-    kv_cache(:,:n_seq2,:,:) = kv_cache2(:,:,:,:)
+    do i4 = 1, m%n_layer
+    do i3 = 1, 2
+    do i2 = 1, n_seq2
+    do i1 = 1, m%n_embd
+        kv_cache(i1,i2,i3,i4) = kv_cache2(i1,i2,i3,i4)
+    end do
+    end do
+    end do
+    end do
     deallocate(kv_cache2)
     next_id = maxloc(logits(:,n_seq_x), dim=1)-1
     input2(n_seq2+1) = next_id

From a27eb59aeb5811e37c7873e5ca9e884f23980b61 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 16:32:35 -0600
Subject: [PATCH 63/84] Print the output as text

---
 driver.f90 | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/driver.f90 b/driver.f90
index 688b5ea..85331d8 100644
--- a/driver.f90
+++ b/driver.f90
@@ -184,10 +184,10 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print *
 print "(a)", "Output tokens:"
 print *, output
-!output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
-!print *
-!print "(a)", "Decoded output as text:"
-!print "(a)", output_txt
+output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
+print *
+print "(a)", "Decoded output as text:"
+print *, output_txt
 if (output(1) /= 703) error stop
 if (output(2) /= 484) error stop
 if (output(19) /= 1517) error stop

From c77237c78ecc237894b2de8dac1a2019ae793ec6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 17:56:59 -0600
Subject: [PATCH 64/84] Remove checks, they are tested in tests

---
 driver.f90 | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/driver.f90 b/driver.f90
index 85331d8..0c5c6c8 100644
--- a/driver.f90
+++ b/driver.f90
@@ -147,11 +147,6 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 !print "(1000(i6))", input
 print *, input
 print *
-if (n_seq /= 19) error stop
-if (input(1) /= 36235) error stop
-if (input(2) /= 39141) error stop
-if (input(18) /= 407) error stop
-if (input(19) /= 5967) error stop
 
 if (n_seq + n_tokens_to_generate >= m%n_ctx) then
     print *, "The maximum sequence length of the model was surpassed."
@@ -188,10 +183,6 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print *
 print "(a)", "Decoded output as text:"
 print *, output_txt
-if (output(1) /= 703) error stop
-if (output(2) /= 484) error stop
-if (output(19) /= 1517) error stop
-if (output(20) /= 318) error stop
 end subroutine
 
 subroutine gpt2_driver3(input_txt, n_tokens_to_generate, stop_text, m, output_txt)

From da839e63514c36180764b741bc72c49d4ea83338 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 17:57:24 -0600
Subject: [PATCH 65/84] Remove trailing whitespace (not used by GF)

---
 tests/test_more_inputs.f90 | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_more_inputs.f90 b/tests/test_more_inputs.f90
index 7f3a586..80d5f35 100644
--- a/tests/test_more_inputs.f90
+++ b/tests/test_more_inputs.f90
@@ -11,19 +11,19 @@ program test_more_inputs
 
 call load_model("model.dat", m)
 
-call gpt2_driver2("Ondřej Čertík was born in ", 13, m, input, output)
+call gpt2_driver2("Ondřej Čertík was born in", 13, m, input, output)
 print *
 print *, "TESTS:"
 call test(input, input_ref, "Input")
 call test(output, output_ref, "Output")
 
-call gpt2_driver2("San Francisco is ", 8, m, input, output)
+call gpt2_driver2("San Francisco is", 8, m, input, output)
 print *
 print *, "TESTS:"
 call test(input, [15017, 6033, 318], "Input")
 call test(output, [257, 1748, 286, 517, 621, 352, 1510, 661], "Output")
 
-call gpt2_driver2("Cars are ", 13, m, input, output)
+call gpt2_driver2("Cars are", 13, m, input, output)
 print *
 print *, "TESTS:"
 call test(input, [34, 945, 389], "Input")

From 7e732dd25c80397cec89f608161c375c507d269a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 18:07:19 -0600
Subject: [PATCH 66/84] Printing

---
 driver.f90 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/driver.f90 b/driver.f90
index 0c5c6c8..0133c19 100644
--- a/driver.f90
+++ b/driver.f90
@@ -175,7 +175,7 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 print *
 t2o = omp_get_wtime()
 call cpu_time(t2)
-print "(a,f8.3,a,f4.2,a)", "    done. Time:", t2o-t1o, "s (", (t2-t1)/(t2o-t1o), "x)"
+print "(a,f8.3,a,f8.2,a)", "    done. Time:", t2o-t1o, "s (", (t2-t1)/(t2o-t1o), "x)"
 print *
 print "(a)", "Output tokens:"
 print *, output

From dc02c77e38f140787f06b98f4ccfb2a01998a1f6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 18:07:25 -0600
Subject: [PATCH 67/84] Enable checks

---
 driver.f90 | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/driver.f90 b/driver.f90
index 0133c19..22e39d1 100644
--- a/driver.f90
+++ b/driver.f90
@@ -154,16 +154,16 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
     error stop
 end if
 
-!print "(a)", "Decoded input as text:"
+print "(a)", "Decoded input as text:"
 !print "(a)", decode(input, decoder_idx, decoder_txt, byte_decoder)
-!allocate(character(0) :: output_txt) ! Fix GFortran warning
-!output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder)
-!print "(a)", output_txt
+allocate(character(0) :: output_txt) ! Fix GFortran warning
+output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder)
+print *, output_txt
 print *
 
-!if (input_txt /= output_txt) then
-!    error stop "The decoded input text does not agree with the input text"
-!end if
+if (input_txt /= output_txt) then
+    error stop "The decoded input text does not agree with the input text"
+end if
 
 allocate(output(n_tokens_to_generate))
 print "(a)", "Running model..."

From 51b6a1616f58058c8afcf76e45c4a4f7ae0ada5d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 18:11:58 -0600
Subject: [PATCH 68/84] Enable printing tokens

---
 gpt2.f90 | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index e1db77d..db2d998 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -359,9 +359,9 @@ subroutine generate(output, n_tokens_to_generate, m, &
     deallocate(kv_cache2)
     next_id = maxloc(logits(:,n_seq_x), dim=1)-1
     input2(n_seq2+1) = next_id
-    !last_token = decode([next_id], m%decoder_idx, &
-    !    m%decoder_txt, byte_decoder)
-    !write(*, fmt="(a)", advance="no") last_token
+    last_token = decode([next_id], m%decoder_idx, &
+        m%decoder_txt, byte_decoder)
+    write(*, fmt="(a)", advance="no") last_token
     if (present(stop_text)) then
         output_txt = output_txt // last_token
         if (output_txt(len(output_txt)-len(stop_text)+1:len(output_txt)) == stop_text) then

From 2e9371d3d2566f52654b33b7ff09b4528d7b21ed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 20:05:24 -0700
Subject: [PATCH 69/84] Fix output copying

---
 gpt2.f90 | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/gpt2.f90 b/gpt2.f90
index db2d998..a23b9c2 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -302,7 +302,7 @@ subroutine generate(output, n_tokens_to_generate, m, &
 logical, intent(in) :: use_cache
 integer, intent(in) :: byte_decoder(:)
 character(*), intent(in), optional :: stop_text ! Stop if you see this text
-integer, intent(out) :: output(:)
+integer, allocatable, intent(out) :: output(:)
 real(sp), allocatable :: logits(:,:)
 integer :: i, i1, i2, i3, i4
 integer :: n_seq2, n_seq_x
@@ -370,7 +370,9 @@ subroutine generate(output, n_tokens_to_generate, m, &
     end if
     deallocate(logits)
 end do
-do i = 1, n_tokens_to_generate
+! output = input2(n_seq+1:n_seq2+1)
+allocate(output(n_seq2 - n_seq + 1))
+do i = 1, n_seq2 - n_seq + 1
     output(i) = input2(n_seq+i)
 end do
 end subroutine

From ab49255bd8b6c78e83a9e5319be61b49a2da351d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 20:41:25 -0700
Subject: [PATCH 70/84] Better error message

---
 tokenizer.f90 | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tokenizer.f90 b/tokenizer.f90
index 775777c..8da8e72 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -256,7 +256,10 @@ function decode(tokens, idx, decoder_txt, byte_decoder) result(output)
     ! However for GPT-2 it seems only range 0-323 is used from UTF-32.
     c = utf8_to_codepoint(output2, i)
     ! [0,324] -> [0,255]
-    if (c < 0 .or. c > ubound(byte_decoder,1)) error stop "Codepoint out of range for byte decoder"
+    if (c < 0 .or. c > ubound(byte_decoder,1)) then
+        print *, "Codepoint out of range for byte decoder:", c, ubound(byte_decoder,1)
+        error stop "Codepoint out of range for byte decoder"
+    end if
     tmp = achar(byte_decoder(c))
     output = output // tmp
     if (i == len(output2)) exit

From 547667485984cb87b8acea5f8d757dc9c7746b62 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 21:23:03 -0700
Subject: [PATCH 71/84] Workaround for
 https://github.com/lfortran/lfortran/issues/2069

---
 tokenizer.f90 | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tokenizer.f90 b/tokenizer.f90
index 8da8e72..e66bfb2 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -258,10 +258,12 @@ function decode(tokens, idx, decoder_txt, byte_decoder) result(output)
     ! [0,324] -> [0,255]
     if (c < 0 .or. c > ubound(byte_decoder,1)) then
         print *, "Codepoint out of range for byte decoder:", c, ubound(byte_decoder,1)
-        error stop "Codepoint out of range for byte decoder"
+        ! We skip this, to allow LFortran to run
+!        error stop "Codepoint out of range for byte decoder"
+    else
+        tmp = achar(byte_decoder(c))
+        output = output // tmp
     end if
-    tmp = achar(byte_decoder(c))
-    output = output // tmp
     if (i == len(output2)) exit
     i = i + 1
 end do

From 5f96d8cfe1159861384d46c4bb0d2104f8a84933 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= <ondrej@certik.us>
Date: Mon, 24 Jul 2023 21:43:02 -0700
Subject: [PATCH 72/84] Make it more explicit

---
 tokenizer.f90 | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/tokenizer.f90 b/tokenizer.f90
index e66bfb2..f546c20 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -242,13 +242,25 @@ function decode(tokens, idx, decoder_txt, byte_decoder) result(output)
 integer(1), intent(in) :: decoder_txt(:)
 character(:), allocatable :: output
 character(:), allocatable :: output2, tmp
-integer :: i, c
-allocate(character(0) :: output2) ! Fix GFortran warning
-output2 = ""
+integer, parameter :: max_len = 4096
+integer(1) :: output3(max_len)
+integer(1), allocatable :: output4(:)
+integer :: i, j, c, pos
+pos = 0
 do i = 1, size(tokens)
     if (tokens(i) < 0) error stop "tokens(i) < 0"
-    output2 = output2 // c2s(decoder_txt(idx(tokens(i))+1:idx(tokens(i)+1)))
+    ! output2 = output2 // (decoder_txt(idx(tokens(i))+1:idx(tokens(i)+1)))
+    do j = idx(tokens(i))+1, idx(tokens(i)+1)
+        pos = pos + 1
+        output3(pos) = decoder_txt(j)
+    end do
+end do
+allocate(character(0) :: output2) ! Fix GFortran warning
+allocate(output4(pos))
+do j = 1, pos
+    output4(j) = output3(j)
 end do
+output2 = c2s(output4)
 i = 1
 output = ""
 do
@@ -259,7 +271,7 @@ function decode(tokens, idx, decoder_txt, byte_decoder) result(output)
     if (c < 0 .or. c > ubound(byte_decoder,1)) then
         print *, "Codepoint out of range for byte decoder:", c, ubound(byte_decoder,1)
         ! We skip this, to allow LFortran to run
-!        error stop "Codepoint out of range for byte decoder"
+        error stop "Codepoint out of range for byte decoder"
     else
         tmp = achar(byte_decoder(c))
         output = output // tmp

From cefc3cfc9b9fd0ca56e981c5a309f9b072236a65 Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Tue, 1 Aug 2023 18:27:25 -0700
Subject: [PATCH 73/84] begin-port

---
 driver.f90                  | 502 ++++++++++++-------------
 gpt2.f90                    | 712 ++++++++++++++++++------------------
 lp-fastgpt/README.md        |   0
 lp-fastgpt/create_model.py  | 355 ++++++++++++++++++
 lp-fastgpt/main.py          |  21 ++
 lp-fastgpt/requirements.txt |   0
 main.f90                    |  10 +-
 7 files changed, 988 insertions(+), 612 deletions(-)
 create mode 100644 lp-fastgpt/README.md
 create mode 100644 lp-fastgpt/create_model.py
 create mode 100644 lp-fastgpt/main.py
 create mode 100644 lp-fastgpt/requirements.txt

diff --git a/driver.f90 b/driver.f90
index 22e39d1..c4bf9f0 100644
--- a/driver.f90
+++ b/driver.f90
@@ -1,273 +1,273 @@
 module driver
-use gpt2_mod, only: generate, model_t
-use tokenizer, only: encode, decode, string
-use omp, only: omp_get_wtime
-implicit none
+    use gpt2_mod, only : generate, model_t
+    use tokenizer, only : encode, decode, string
+    use omp, only : omp_get_wtime
+    implicit none
 
-integer, parameter :: sp = kind(0.0)
-integer, parameter :: dp = kind(0.d0)
-character(1), parameter :: LF = achar(10)
+    integer, parameter :: sp = kind(0.0)
+    integer, parameter :: dp = kind(0.d0)
+    character(1), parameter :: LF = achar(10)
 
 contains
 
-subroutine load_input(filename, input_txt, n_tokens_to_generate)
-! Load the input from a namelist `filename`
-character(*), intent(in) :: filename
-character(:), allocatable, intent(out) :: input_txt
-integer, intent(out) :: n_tokens_to_generate
-character(1024) :: input_txt2
-integer :: u, ios
-!namelist / input_fastGPT / n_tokens_to_generate
-allocate(character(0) :: input_txt)
-n_tokens_to_generate = 20
-input_txt = "Alan Turing theorized that computers would one day become very powerful, but even he could not imagine"
-!open(newunit=u, file=filename, status="old")
-!read(u, input_fastGPT)
-!do
-!    read(u, "(a)", iostat=ios) input_txt2
-!    if (ios /= 0) exit
-!    if (len(input_txt) > 0) input_txt = input_txt // char(10)
-!    input_txt = input_txt // trim(input_txt2)
-!end do
-!close(u)
-end subroutine
+    subroutine load_input(filename, input_txt, n_tokens_to_generate)
+        ! Load the input from a namelist `filename`
+        character(*), intent(in) :: filename
+        character(:), allocatable, intent(out) :: input_txt
+        integer, intent(out) :: n_tokens_to_generate
+        character(1024) :: input_txt2
+        integer :: u, ios
+        !namelist / input_fastGPT / n_tokens_to_generate
+        allocate(character(0) :: input_txt)
+        n_tokens_to_generate = 20
+        input_txt = "Alan Turing theorized that computers would one day become very powerful, but even he could not imagine"
+        !open(newunit=u, file=filename, status="old")
+        !read(u, input_fastGPT)
+        !do
+        !    read(u, "(a)", iostat=ios) input_txt2
+        !    if (ios /= 0) exit
+        !    if (len(input_txt) > 0) input_txt = input_txt // char(10)
+        !    input_txt = input_txt // trim(input_txt2)
+        !end do
+        !close(u)
+    end subroutine
 
-subroutine load_model(filename, m)
-character(*), intent(in) :: filename
-type(model_t), intent(out) :: m
-! We use the following fastGPT model type number
-!   fastGPT (digits look similar to the letters they represent)
-! 0xfa51697 = 262477463
-integer, parameter :: current_model_mark = 262477463
-integer, parameter :: current_model_version = 1
-integer :: model_mark
-integer :: u
-open(newunit=u, file=filename, form="unformatted", access="stream", status="old")
-read(u) model_mark
-if (model_mark /= current_model_mark) then
-    print *, "Found:", model_mark
-    print *, "Expected:", current_model_mark
-    error stop "Invalid fastGPT model file"
-end if
-read(u) m%model_file_version
-if (m%model_file_version /= current_model_version) then
-    print *, "Found:", m%model_file_version
-    print *, "Expected:", current_model_version
-    error stop "Incompatible model version"
-end if
-read(u) m%n_vocab, m%n_ctx, m%n_embd, m%n_layer, m%n_head, m%n_decoder_idx, &
-    m%n_decoder_txt, m%n_vocab_idx, m%n_vocab_txt, m%n_byte_encoder
-allocate(m%wte(m%n_embd,m%n_vocab), m%wpe(m%n_embd,m%n_ctx), &
-    m%mlp_fc_w(4*m%n_embd,m%n_embd,m%n_layer), m%mlp_fc_b(4*m%n_embd,m%n_layer), &
-    m%mlp_proj_w(m%n_embd,4*m%n_embd,m%n_layer), m%mlp_proj_b(m%n_embd,m%n_layer), &
-    m%attn_w(3*m%n_embd,m%n_embd,m%n_layer), m%attn_b(3*m%n_embd,m%n_layer), &
-    m%attn_proj_w(m%n_embd,m%n_embd,m%n_layer), m%attn_proj_b(m%n_embd,m%n_layer), &
-    m%ln1_b(m%n_embd,m%n_layer), m%ln1_g(m%n_embd,m%n_layer), &
-    m%ln2_b(m%n_embd,m%n_layer), m%ln2_g(m%n_embd,m%n_layer), &
-    m%lnf_b(m%n_embd), m%lnf_g(m%n_embd), &
-    m%decoder_idx(0:m%n_decoder_idx-1), m%decoder_txt(m%n_decoder_txt), &
-    m%vocab_idx(0:m%n_vocab_idx-1), m%vocab_txt(m%n_vocab_txt), &
-    m%byte_encoder(0:m%n_byte_encoder-1))
-read(u) m%wte, m%wpe, &
-    m%mlp_fc_w, m%mlp_fc_b, &
-    m%mlp_proj_w, m%mlp_proj_b, &
-    m%attn_w, m%attn_b, &
-    m%attn_proj_w, m%attn_proj_b, &
-    m%ln1_b, m%ln1_g, &
-    m%ln2_b, m%ln2_g, &
-    m%lnf_b, m%lnf_g, &
-    m%decoder_idx, m%decoder_txt, &
-    m%vocab_idx, m%vocab_txt, &
-    m%byte_encoder
-close(u)
-end subroutine
+    subroutine load_model(filename, m)
+        character(*), intent(in) :: filename
+        type(model_t), intent(out) :: m
+        ! We use the following fastGPT model type number
+        !   fastGPT (digits look similar to the letters they represent)
+        ! 0xfa51697 = 262477463
+        integer, parameter :: current_model_mark = 262477463
+        integer, parameter :: current_model_version = 1
+        integer :: model_mark
+        integer :: u
+        open(newunit = u, file = filename, form = "unformatted", access = "stream", status = "old")
+        read(u) model_mark
+        if (model_mark /= current_model_mark) then
+            print *, "Found:", model_mark
+            print *, "Expected:", current_model_mark
+            error stop "Invalid fastGPT model file"
+        end if
+        read(u) m%model_file_version
+        if (m%model_file_version /= current_model_version) then
+            print *, "Found:", m%model_file_version
+            print *, "Expected:", current_model_version
+            error stop "Incompatible model version"
+        end if
+        read(u) m%n_vocab, m%n_ctx, m%n_embd, m%n_layer, m%n_head, m%n_decoder_idx, &
+                m%n_decoder_txt, m%n_vocab_idx, m%n_vocab_txt, m%n_byte_encoder
+        allocate(m%wte(m%n_embd, m%n_vocab), m%wpe(m%n_embd, m%n_ctx), &
+                m%mlp_fc_w(4 * m%n_embd, m%n_embd, m%n_layer), m%mlp_fc_b(4 * m%n_embd, m%n_layer), &
+                m%mlp_proj_w(m%n_embd, 4 * m%n_embd, m%n_layer), m%mlp_proj_b(m%n_embd, m%n_layer), &
+                m%attn_w(3 * m%n_embd, m%n_embd, m%n_layer), m%attn_b(3 * m%n_embd, m%n_layer), &
+                m%attn_proj_w(m%n_embd, m%n_embd, m%n_layer), m%attn_proj_b(m%n_embd, m%n_layer), &
+                m%ln1_b(m%n_embd, m%n_layer), m%ln1_g(m%n_embd, m%n_layer), &
+                m%ln2_b(m%n_embd, m%n_layer), m%ln2_g(m%n_embd, m%n_layer), &
+                m%lnf_b(m%n_embd), m%lnf_g(m%n_embd), &
+                m%decoder_idx(0:m%n_decoder_idx - 1), m%decoder_txt(m%n_decoder_txt), &
+                m%vocab_idx(0:m%n_vocab_idx - 1), m%vocab_txt(m%n_vocab_txt), &
+                m%byte_encoder(0:m%n_byte_encoder - 1))
+        read(u) m%wte, m%wpe, &
+                m%mlp_fc_w, m%mlp_fc_b, &
+                m%mlp_proj_w, m%mlp_proj_b, &
+                m%attn_w, m%attn_b, &
+                m%attn_proj_w, m%attn_proj_b, &
+                m%ln1_b, m%ln1_g, &
+                m%ln2_b, m%ln2_g, &
+                m%lnf_b, m%lnf_g, &
+                m%decoder_idx, m%decoder_txt, &
+                m%vocab_idx, m%vocab_txt, &
+                m%byte_encoder
+        close(u)
+    end subroutine
 
-subroutine gpt2_driver(input, output, m)
-integer, allocatable, intent(out) :: input(:), output(:)
-type(model_t), intent(out) :: m
-character(:), allocatable :: input_txt
-integer :: n_tokens_to_generate
-real(dp) :: t1, t2
-call load_input("input", input_txt, n_tokens_to_generate)
+    subroutine gpt2_driver(input, output, m)
+        integer, allocatable, intent(out) :: input(:), output(:)
+        type(model_t), intent(out) :: m
+        character(:), allocatable :: input_txt
+        integer :: n_tokens_to_generate
+        real(dp) :: t1, t2
+        call load_input("input", input_txt, n_tokens_to_generate)
 
-! Load the model
-print "(a)", "Loading the model..."
-call cpu_time(t1)
-call load_model("model.dat", m)
-call cpu_time(t2)
-print "(a,f8.3,a,i2)", "    done. Time:", t2-t1, "s, Model file version:", m%model_file_version
-print *
-print "(a)", "Model parameters:"
-print "(a,i6)", "n_vocab =", m%n_vocab
-print "(a,i6)", "n_ctx   =", m%n_ctx
-print "(a,i6)", "n_embd  =", m%n_embd
-print "(a,i6)", "n_layer =", m%n_layer
-print "(a,i6)", "n_head  =", m%n_head
-print *
+        ! Load the model
+        print "(a)", "Loading the model..."
+        call cpu_time(t1)
+        call load_model("model.dat", m)
+        call cpu_time(t2)
+        print "(a,f8.3,a,i2)", "    done. Time:", t2 - t1, "s, Model file version:", m%model_file_version
+        print *
+        print "(a)", "Model parameters:"
+        print "(a,i6)", "n_vocab =", m%n_vocab
+        print "(a,i6)", "n_ctx   =", m%n_ctx
+        print "(a,i6)", "n_embd  =", m%n_embd
+        print "(a,i6)", "n_layer =", m%n_layer
+        print "(a,i6)", "n_head  =", m%n_head
+        print *
 
-call gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
-endsubroutine
+        call gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
+    endsubroutine
 
-subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
-character(*), intent(in) :: input_txt
-integer, intent(in) :: n_tokens_to_generate
-type(model_t), intent(in) :: m
-integer, allocatable, intent(out) :: input(:), output(:)
-integer, allocatable :: byte_decoder(:)
-integer :: n_seq
-character(:), allocatable :: output_txt
-real(dp) :: t1, t2, t1o, t2o
-integer :: i
-logical :: use_cache
+    subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
+        character(*), intent(in) :: input_txt
+        integer, intent(in) :: n_tokens_to_generate
+        type(model_t), intent(in) :: m
+        integer, allocatable, intent(out) :: input(:), output(:)
+        integer, allocatable :: byte_decoder(:)
+        integer :: n_seq
+        character(:), allocatable :: output_txt
+        real(dp) :: t1, t2, t1o, t2o
+        integer :: i
+        logical :: use_cache
 
-! Compute byte_decoder:
-allocate(byte_decoder(0:maxval(m%byte_encoder)))
-byte_decoder = 0
-do i = 0, size(m%byte_encoder)-1
-    byte_decoder(m%byte_encoder(i)) = i
-end do
+        ! Compute byte_decoder:
+        allocate(byte_decoder(0:maxval(m%byte_encoder)))
+        byte_decoder = 0
+        do i = 0, size(m%byte_encoder) - 1
+            byte_decoder(m%byte_encoder(i)) = i
+        end do
 
-print "(a)", "Input text"
-print "(a)", input_txt
+        print "(a)", "Input text"
+        print "(a)", input_txt
 
-print *
-print "(a)",  "Encoding: tokenizing input text into tokens (currently slow)..."
-call cpu_time(t1)
-input = encode(input_txt, m%decoder_idx, m%decoder_txt, m%vocab_idx, m%vocab_txt, &
-    m%byte_encoder)
-call cpu_time(t2)
-n_seq = size(input)
-print "(a,f8.3,a)", "    done. Time:", t2-t1, "s"
-print *
-print "(a)", "Input parameters:"
-print "(a,i4)", "n_seq                =", n_seq
-print "(a,i4)", "n_tokens_to_generate =", n_tokens_to_generate
-print *
-print "(a)", "Input tokens:"
-!print "(1000(i6))", input
-print *, input
-print *
+        print *
+        print "(a)", "Encoding: tokenizing input text into tokens (currently slow)..."
+        call cpu_time(t1)
+        input = encode(input_txt, m%decoder_idx, m%decoder_txt, m%vocab_idx, m%vocab_txt, &
+                m%byte_encoder)
+        call cpu_time(t2)
+        n_seq = size(input)
+        print "(a,f8.3,a)", "    done. Time:", t2 - t1, "s"
+        print *
+        print "(a)", "Input parameters:"
+        print "(a,i4)", "n_seq                =", n_seq
+        print "(a,i4)", "n_tokens_to_generate =", n_tokens_to_generate
+        print *
+        print "(a)", "Input tokens:"
+        !print "(1000(i6))", input
+        print *, input
+        print *
 
-if (n_seq + n_tokens_to_generate >= m%n_ctx) then
-    print *, "The maximum sequence length of the model was surpassed."
-    print *, "Make the input and/or number of tokens to generate shorter."
-    error stop
-end if
+        if (n_seq + n_tokens_to_generate >= m%n_ctx) then
+            print *, "The maximum sequence length of the model was surpassed."
+            print *, "Make the input and/or number of tokens to generate shorter."
+            error stop
+        end if
 
-print "(a)", "Decoded input as text:"
-!print "(a)", decode(input, decoder_idx, decoder_txt, byte_decoder)
-allocate(character(0) :: output_txt) ! Fix GFortran warning
-output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder)
-print *, output_txt
-print *
+        print "(a)", "Decoded input as text:"
+        !print "(a)", decode(input, decoder_idx, decoder_txt, byte_decoder)
+        allocate(character(0) :: output_txt) ! Fix GFortran warning
+        output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder)
+        print *, output_txt
+        print *
 
-if (input_txt /= output_txt) then
-    error stop "The decoded input text does not agree with the input text"
-end if
+        if (input_txt /= output_txt) then
+            error stop "The decoded input text does not agree with the input text"
+        end if
 
-allocate(output(n_tokens_to_generate))
-print "(a)", "Running model..."
-call cpu_time(t1)
-t1o = omp_get_wtime()
-use_cache = .true.
-call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
-    byte_decoder)
-print *
-t2o = omp_get_wtime()
-call cpu_time(t2)
-print "(a,f8.3,a,f8.2,a)", "    done. Time:", t2o-t1o, "s (", (t2-t1)/(t2o-t1o), "x)"
-print *
-print "(a)", "Output tokens:"
-print *, output
-output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
-print *
-print "(a)", "Decoded output as text:"
-print *, output_txt
-end subroutine
+        allocate(output(n_tokens_to_generate))
+        print "(a)", "Running model..."
+        call cpu_time(t1)
+        t1o = omp_get_wtime()
+        use_cache = .true.
+        call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
+                byte_decoder)
+        print *
+        t2o = omp_get_wtime()
+        call cpu_time(t2)
+        print "(a,f8.3,a,f8.2,a)", "    done. Time:", t2o - t1o, "s (", (t2 - t1) / (t2o - t1o), "x)"
+        print *
+        print "(a)", "Output tokens:"
+        print *, output
+        output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
+        print *
+        print "(a)", "Decoded output as text:"
+        print *, output_txt
+    end subroutine
 
-subroutine gpt2_driver3(input_txt, n_tokens_to_generate, stop_text, m, output_txt)
-character(*), intent(in) :: input_txt, stop_text
-integer, intent(in) :: n_tokens_to_generate
-type(model_t), intent(in) :: m
-integer, allocatable :: input(:), output(:)
-integer, allocatable :: byte_decoder(:)
-integer :: n_seq
-character(:), allocatable, intent(out) :: output_txt
-integer :: i
-logical :: use_cache
-! TODO: move the decoder into model_t
-! Compute byte_decoder:
-allocate(byte_decoder(0:maxval(m%byte_encoder)))
-byte_decoder = 0
-do i = 0, size(m%byte_encoder)-1
-    byte_decoder(m%byte_encoder(i)) = i
-end do
-input = encode(input_txt, m%decoder_idx, m%decoder_txt, m%vocab_idx, m%vocab_txt, &
-    m%byte_encoder)
-n_seq = size(input)
-if (n_seq + n_tokens_to_generate >= m%n_ctx) then
-    print *, "The maximum sequence length of the model was surpassed."
-    print *, "Make the input and/or number of tokens to generate shorter."
-    error stop
-end if
-allocate(character(0) :: output_txt) ! Fix GFortran warning
-output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder)
-if (input_txt /= output_txt) then
-    error stop "The decoded input text does not agree with the input text"
-end if
-use_cache = .true.
-call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
-    byte_decoder, stop_text)
-output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
-end subroutine
+    subroutine gpt2_driver3(input_txt, n_tokens_to_generate, stop_text, m, output_txt)
+        character(*), intent(in) :: input_txt, stop_text
+        integer, intent(in) :: n_tokens_to_generate
+        type(model_t), intent(in) :: m
+        integer, allocatable :: input(:), output(:)
+        integer, allocatable :: byte_decoder(:)
+        integer :: n_seq
+        character(:), allocatable, intent(out) :: output_txt
+        integer :: i
+        logical :: use_cache
+        ! TODO: move the decoder into model_t
+        ! Compute byte_decoder:
+        allocate(byte_decoder(0:maxval(m%byte_encoder)))
+        byte_decoder = 0
+        do i = 0, size(m%byte_encoder) - 1
+            byte_decoder(m%byte_encoder(i)) = i
+        end do
+        input = encode(input_txt, m%decoder_idx, m%decoder_txt, m%vocab_idx, m%vocab_txt, &
+                m%byte_encoder)
+        n_seq = size(input)
+        if (n_seq + n_tokens_to_generate >= m%n_ctx) then
+            print *, "The maximum sequence length of the model was surpassed."
+            print *, "Make the input and/or number of tokens to generate shorter."
+            error stop
+        end if
+        allocate(character(0) :: output_txt) ! Fix GFortran warning
+        output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder)
+        if (input_txt /= output_txt) then
+            error stop "The decoded input text does not agree with the input text"
+        end if
+        use_cache = .true.
+        call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, &
+                byte_decoder, stop_text)
+        output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder)
+    end subroutine
 
-function get_prompt() result(input)
-character(:), allocatable :: input
-character(1024) :: tmp
-integer ::ios
-read(*,"(a)",iostat=ios) tmp
-if (ios == 0) then
-    input = trim(tmp)
-else
-    input = ""
-end if
-end function
+    function get_prompt() result(input)
+        character(:), allocatable :: input
+        character(1024) :: tmp
+        integer :: ios
+        read(*, "(a)", iostat = ios) tmp
+        if (ios == 0) then
+            input = trim(tmp)
+        else
+            input = ""
+        end if
+    end function
 
-subroutine chat(present_inputs, inputs)
-logical, intent(in) :: present_inputs
-type(string), optional, intent(in) :: inputs(:)
-type(model_t) :: m
-character(:), allocatable :: prompt, input, output
-integer :: i, n_prompts
-call load_model("model.dat", m)
-prompt = "Your name is fastGPT and you are an AI bot. The user will ask you &
-&questions and you answer in a nice, truthful, short way." // LF // "&
-&User: What is the capital of Czechia?" // LF // "&
-&fastGPT: Prague." // LF // "&
-&User: How many legs does a dog have?" // LF // "&
-&fastGPT: Four." // LF // "&
-&User:"
-write(*,"(a)",advance="no") prompt
-if (present_inputs) then
-    n_prompts = size(inputs)
-else
-    n_prompts = 1024
-end if
-do i = 1, n_prompts
-    write(*,"(a)",advance="no")  " "
-    if (present_inputs) then
-        input = inputs(i)%s
-        write(*,"(a)") input
-    else
-        input = get_prompt()
-        if (input == "") exit
-    end if
-    write(*,"(a)",advance="no") "fastGPT:"
-    prompt = prompt // " " // input // LF // "fastGPT:"
-    call gpt2_driver3(prompt, 200, "User:", m, output)
-    prompt = prompt // output
-end do
-print *
-end subroutine
+    subroutine chat(present_inputs, inputs)
+        logical, intent(in) :: present_inputs
+        type(string), optional, intent(in) :: inputs(:)
+        type(model_t) :: m
+        character(:), allocatable :: prompt, input, output
+        integer :: i, n_prompts
+        call load_model("model.dat", m)
+        prompt = "Your name is fastGPT and you are an AI bot. The user will ask you &
+                &questions and you answer in a nice, truthful, short way." // LF // "&
+                &User: What is the capital of Czechia?" // LF // "&
+                &fastGPT: Prague." // LF // "&
+                &User: How many legs does a dog have?" // LF // "&
+                &fastGPT: Four." // LF // "&
+                &User:"
+        write(*, "(a)", advance = "no") prompt
+        if (present_inputs) then
+            n_prompts = size(inputs)
+        else
+            n_prompts = 1024
+        end if
+        do i = 1, n_prompts
+            write(*, "(a)", advance = "no")  " "
+            if (present_inputs) then
+                input = inputs(i)%s
+                write(*, "(a)") input
+            else
+                input = get_prompt()
+                if (input == "") exit
+            end if
+            write(*, "(a)", advance = "no") "fastGPT:"
+            prompt = prompt // " " // input // LF // "fastGPT:"
+            call gpt2_driver3(prompt, 200, "User:", m, output)
+            prompt = prompt // output
+        end do
+        print *
+    end subroutine
 
 end module
diff --git a/gpt2.f90 b/gpt2.f90
index a23b9c2..731259b 100644
--- a/gpt2.f90
+++ b/gpt2.f90
@@ -1,380 +1,380 @@
 module gpt2_mod
-use linalg, only: matmul_2d, matmul_2d_t
-use tokenizer, only: decode
-implicit none
+    use linalg, only : matmul_2d, matmul_2d_t
+    use tokenizer, only : decode
+    implicit none
 
-integer, parameter :: sp = kind(0.0)
-real(sp), parameter :: pi = 3.14159265358979323846_sp
+    integer, parameter :: sp = kind(0.0)
+    real(sp), parameter :: pi = 3.14159265358979323846_sp
 
-! This derived type contains all the data of the GPT-2 model, including all
-! weights, model parameters, and encoder/decoder data
-type :: model_t
-    integer :: n_vocab, n_ctx, n_embd, n_layer, n_head, &
-        n_decoder_idx, n_decoder_txt, &
-        n_vocab_idx, n_vocab_txt, n_byte_encoder
-    real(sp), allocatable :: wte(:,:), wpe(:,:), &
-        mlp_fc_w(:,:,:), mlp_fc_b(:,:), &
-        mlp_proj_w(:,:,:), mlp_proj_b(:,:), &
-        attn_w(:,:,:), attn_b(:,:), &
-        attn_proj_w(:,:,:), attn_proj_b(:,:), &
-        ln1_b(:,:), ln1_g(:,:), &
-        ln2_b(:,:), ln2_g(:,:), &
-        lnf_b(:), lnf_g(:)
-    integer, allocatable :: decoder_idx(:), vocab_idx(:), byte_encoder(:)
-    integer(1), allocatable :: decoder_txt(:), vocab_txt(:)
-    integer :: model_file_version
-end type
+    ! This derived type contains all the data of the GPT-2 model, including all
+    ! weights, model parameters, and encoder/decoder data
+    type :: model_t
+        integer :: n_vocab, n_ctx, n_embd, n_layer, n_head, &
+                n_decoder_idx, n_decoder_txt, &
+                n_vocab_idx, n_vocab_txt, n_byte_encoder
+        real(sp), allocatable :: wte(:, :), wpe(:, :), &
+                mlp_fc_w(:, :, :), mlp_fc_b(:, :), &
+                mlp_proj_w(:, :, :), mlp_proj_b(:, :), &
+                attn_w(:, :, :), attn_b(:, :), &
+                attn_proj_w(:, :, :), attn_proj_b(:, :), &
+                ln1_b(:, :), ln1_g(:, :), &
+                ln2_b(:, :), ln2_g(:, :), &
+                lnf_b(:), lnf_g(:)
+        integer, allocatable :: decoder_idx(:), vocab_idx(:), byte_encoder(:)
+        integer(1), allocatable :: decoder_txt(:), vocab_txt(:)
+        integer :: model_file_version
+    end type
 
 contains
 
-elemental real(sp) function fast_tanh(x) result(y)
-real(sp), intent(in) :: x
-real(sp) :: x2
-if (x > 5) then
-    y = 1
-elseif (x < -5) then
-    y = -1
-else
-    x2 = x*x
-    y = x * (0.98569772605911309407 + x2 *(-0.2794500993392901382 &
-        + x2 * (6.8280504526399188164e-2 + x2 * (-1.0972014877337651823e-2 &
-        + x2 * (1.1132367134444316902e-3 + x2 * (-7.018851897305717565e-5 &
-        + x2 * (2.656616768082727089e-6 + x2 * (-5.5138381821615909058e-8 &
-        + x2 * 4.8162484477588665996e-10))))))))
-end if
-end function
+    elemental real(sp) function fast_tanh(x) result(y)
+        real(sp), intent(in) :: x
+        real(sp) :: x2
+        if (x > 5) then
+            y = 1
+        elseif (x < -5) then
+            y = -1
+        else
+            x2 = x * x
+            y = x * (0.98569772605911309407 + x2 * (-0.2794500993392901382 &
+                    + x2 * (6.8280504526399188164e-2 + x2 * (-1.0972014877337651823e-2 &
+                            + x2 * (1.1132367134444316902e-3 + x2 * (-7.018851897305717565e-5 &
+                                    + x2 * (2.656616768082727089e-6 + x2 * (-5.5138381821615909058e-8 &
+                                            + x2 * 4.8162484477588665996e-10))))))))
+        end if
+    end function
 
-elemental real(sp) function gelu(x) result(y)
-real(sp), intent(in) :: x
-y = 0.5_sp * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715_sp * x**3)))
-end function
+    elemental real(sp) function gelu(x) result(y)
+        real(sp), intent(in) :: x
+        y = 0.5_sp * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715_sp * x**3)))
+    end function
 
-function softmax(x) result(y)
-real(sp), intent(in) :: x(:,:)
-real(sp) :: y(size(x,1),size(x,2))
-integer :: i,j
-real(sp) :: s
-do i = 1, size(x,2)
-    s = -1e10
-    do j = 1, size(x,1)
-        if (x(j,i) > s) s = x(j,i)
-    end do
-    do j = 1, size(x,1)
-        y(j,i) = exp(x(j,i) - s)
-    end do
-    s = 0
-    do j = 1, size(x,1)
-        s = s + y(j,i)
-    end do
-    do j = 1, size(x,1)
-        y(j,i) = y(j,i) / s
-    end do
-end do
-end function
+    function softmax(x) result(y)
+        real(sp), intent(in) :: x(:, :)
+        real(sp) :: y(size(x, 1), size(x, 2))
+        integer :: i, j
+        real(sp) :: s
+        do i = 1, size(x, 2)
+            s = -1e10
+            do j = 1, size(x, 1)
+                if (x(j, i) > s) s = x(j, i)
+            end do
+            do j = 1, size(x, 1)
+                y(j, i) = exp(x(j, i) - s)
+            end do
+            s = 0
+            do j = 1, size(x, 1)
+                s = s + y(j, i)
+            end do
+            do j = 1, size(x, 1)
+                y(j, i) = y(j, i) / s
+            end do
+        end do
+    end function
 
-subroutine layer_norm(y, x, g, b, eps)
-real(sp), intent(in) :: x(:,:), g(:), b(:), eps
-real(sp), intent(out) :: y(size(x,1),size(x,2))
-real(sp) :: mean(size(x,2)), variance(size(x,2))
-real(sp) :: xi(size(x,1))
-integer :: i, j
-do i = 1, size(x,2)
-    do j = 1, size(x,1)
-        xi(j) = x(j,i)
-    end do
-    mean(i) = sum(xi) / size(x,1)
-    do j = 1, size(x,1)
-        xi(j) = (xi(j) - mean(i))**2
-    end do
-    variance(i) = sum(xi) / size(x,1)
-end do
-!do i = 1, size(x,1)
-!    y(i,:) = (x(i,:) - mean(:)) / sqrt(variance(:) + eps)
-!    y(i,:) = g(i) * y(i,:) + b(i)
-!end do
-do i = 1, size(x,2)
-    do j = 1, size(x,1)
-        y(j,i) = (x(j,i) - mean(i)) / sqrt(variance(i) + eps)
-        y(j,i) = g(j) * y(j,i) + b(j)
-    end do
-end do
-end subroutine
+    subroutine layer_norm(y, x, g, b, eps)
+        real(sp), intent(in) :: x(:, :), g(:), b(:), eps
+        real(sp), intent(out) :: y(size(x, 1), size(x, 2))
+        real(sp) :: mean(size(x, 2)), variance(size(x, 2))
+        real(sp) :: xi(size(x, 1))
+        integer :: i, j
+        do i = 1, size(x, 2)
+            do j = 1, size(x, 1)
+                xi(j) = x(j, i)
+            end do
+            mean(i) = sum(xi) / size(x, 1)
+            do j = 1, size(x, 1)
+                xi(j) = (xi(j) - mean(i))**2
+            end do
+            variance(i) = sum(xi) / size(x, 1)
+        end do
+        !do i = 1, size(x,1)
+        !    y(i,:) = (x(i,:) - mean(:)) / sqrt(variance(:) + eps)
+        !    y(i,:) = g(i) * y(i,:) + b(i)
+        !end do
+        do i = 1, size(x, 2)
+            do j = 1, size(x, 1)
+                y(j, i) = (x(j, i) - mean(i)) / sqrt(variance(i) + eps)
+                y(j, i) = g(j) * y(j, i) + b(j)
+            end do
+        end do
+    end subroutine
 
-function linear(x, w, b) result(y)
-real(sp), intent(in) :: x(:,:), w(:,:), b(:)
-real(sp) :: y(size(b,1),size(x,2))
-integer :: i, j
-!y = matmul(w, x) + spread(b, 2, size(x,2))
-!y = matmul(w, x)
-call matmul_2d(w, x, y)
-do i = 1, size(y,2)
-    do j = 1, size(y,1)
-        y(j,i) = y(j,i) + b(j)
-    end do
-end do
-end function
+    function linear(x, w, b) result(y)
+        real(sp), intent(in) :: x(:, :), w(:, :), b(:)
+        real(sp) :: y(size(b, 1), size(x, 2))
+        integer :: i, j
+        !y = matmul(w, x) + spread(b, 2, size(x,2))
+        !y = matmul(w, x)
+        call matmul_2d(w, x, y)
+        do i = 1, size(y, 2)
+            do j = 1, size(y, 1)
+                y(j, i) = y(j, i) + b(j)
+            end do
+        end do
+    end function
 
-subroutine ffn(y, x, fc_w, fc_b, proj_w, proj_b)
-real(sp), intent(in) :: x(:,:), fc_w(:,:), fc_b(:), proj_w(:,:), proj_b(:)
-real(sp), intent(inout) :: y(size(x,1),size(x,2))
-real(sp) :: yy(size(x,1),size(x,2))
-real(sp) :: a(4*size(x,1),size(x,2))
-real(sp) :: aa(4*size(x,1),size(x,2))
-aa = linear(x, fc_w, fc_b)
-a = gelu(aa)
-yy = linear(a, proj_w, proj_b)
-y = y + yy
-end subroutine
+    subroutine ffn(y, x, fc_w, fc_b, proj_w, proj_b)
+        real(sp), intent(in) :: x(:, :), fc_w(:, :), fc_b(:), proj_w(:, :), proj_b(:)
+        real(sp), intent(inout) :: y(size(x, 1), size(x, 2))
+        real(sp) :: yy(size(x, 1), size(x, 2))
+        real(sp) :: a(4 * size(x, 1), size(x, 2))
+        real(sp) :: aa(4 * size(x, 1), size(x, 2))
+        aa = linear(x, fc_w, fc_b)
+        a = gelu(aa)
+        yy = linear(a, proj_w, proj_b)
+        y = y + yy
+    end subroutine
 
-subroutine attention(y, n_embd_head,n_seq,n_seq_x, q, k, v, mask)
-integer, intent(in) :: n_embd_head, n_seq, n_seq_x
-real(sp), intent(in) :: q(n_embd_head,n_seq_x), k(n_embd_head,n_seq), v(n_embd_head,n_seq), mask(n_seq,n_seq_x)
-real(sp), intent(out) :: y(n_embd_head,n_seq_x)
-real(sp) :: tmp(n_seq,n_seq_x)
-integer :: i, j
-!tmp = matmul(transpose(k), q)
-!call matmul_2d(transpose(k), q, tmp)
-call matmul_2d_t(k, q, tmp)
-do i = 1, n_seq_x
-do j = 1, n_seq
-    tmp(j,i) = tmp(j,i) / sqrt(real(n_embd_head,sp)) + mask(j,i)
-end do
-end do
-tmp = softmax(tmp)
-call matmul_2d(v, tmp, y)
-end subroutine
+    subroutine attention(y, n_embd_head, n_seq, n_seq_x, q, k, v, mask)
+        integer, intent(in) :: n_embd_head, n_seq, n_seq_x
+        real(sp), intent(in) :: q(n_embd_head, n_seq_x), k(n_embd_head, n_seq), v(n_embd_head, n_seq), mask(n_seq, n_seq_x)
+        real(sp), intent(out) :: y(n_embd_head, n_seq_x)
+        real(sp) :: tmp(n_seq, n_seq_x)
+        integer :: i, j
+        !tmp = matmul(transpose(k), q)
+        !call matmul_2d(transpose(k), q, tmp)
+        call matmul_2d_t(k, q, tmp)
+        do i = 1, n_seq_x
+            do j = 1, n_seq
+                tmp(j, i) = tmp(j, i) / sqrt(real(n_embd_head, sp)) + mask(j, i)
+            end do
+        end do
+        tmp = softmax(tmp)
+        call matmul_2d(v, tmp, y)
+    end subroutine
 
-subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head, &
+    subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head, &
             use_kv_cache, kv_cache)
-integer, intent(in) :: n_seq, n_seq_x, n_embd
-real(sp), intent(in) :: x(n_embd,n_seq_x), &
-    attn_w(3*n_embd,n_embd), attn_b(3*n_embd), &
-    proj_w(n_embd,n_embd), proj_b(n_embd)
-real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
-integer, intent(in) :: n_head
-logical, intent(in) :: use_kv_cache
-real(sp), intent(out) :: y(n_embd,n_seq_x)
-real(sp) :: causal_mask(n_seq,n_seq_x)
-real(sp) :: x2(3*n_embd,n_seq_x)
-real(sp) :: q(n_embd/n_head,n_seq_x), k(n_embd/n_head,n_seq), v(n_embd/n_head,n_seq)
-real(sp) :: yy(n_embd/n_head,n_seq_x)
-integer :: i, j, l
-! Mask
-if (use_kv_cache) then
-    causal_mask = 0
-else
-    do j = 1, n_seq
-    do i = 1, n_seq
-        if (i > j) then
-            causal_mask(i,j) = -1e10_sp
+        integer, intent(in) :: n_seq, n_seq_x, n_embd
+        real(sp), intent(in) :: x(n_embd, n_seq_x), &
+                attn_w(3 * n_embd, n_embd), attn_b(3 * n_embd), &
+                proj_w(n_embd, n_embd), proj_b(n_embd)
+        real(sp), intent(inout) :: kv_cache(n_embd, n_seq, 2)
+        integer, intent(in) :: n_head
+        logical, intent(in) :: use_kv_cache
+        real(sp), intent(out) :: y(n_embd, n_seq_x)
+        real(sp) :: causal_mask(n_seq, n_seq_x)
+        real(sp) :: x2(3 * n_embd, n_seq_x)
+        real(sp) :: q(n_embd / n_head, n_seq_x), k(n_embd / n_head, n_seq), v(n_embd / n_head, n_seq)
+        real(sp) :: yy(n_embd / n_head, n_seq_x)
+        integer :: i, j, l
+        ! Mask
+        if (use_kv_cache) then
+            causal_mask = 0
         else
-            causal_mask(i,j) = 0
+            do j = 1, n_seq
+                do i = 1, n_seq
+                    if (i > j) then
+                        causal_mask(i, j) = -1e10_sp
+                    else
+                        causal_mask(i, j) = 0
+                    end if
+                end do
+            end do
         end if
-    end do
-    end do
-end if
-x2 = linear(x, attn_w, attn_b)
-if (use_kv_cache) then
-    do j = 1, n_embd
-        kv_cache(j,n_seq,1) = x2((2-1)*n_embd+j,1)
-        kv_cache(j,n_seq,2) = x2((3-1)*n_embd+j,1)
-    end do
-else
-    do i = 1, n_seq
-    do j = 1, n_embd
-        kv_cache(j,i,1) = x2((2-1)*n_embd+j,i)
-        kv_cache(j,i,2) = x2((3-1)*n_embd+j,i)
-    end do
-    end do
-end if
-! Perform attention over each head
-do l = 1, n_head
-    do i = 1, n_seq_x
-    do j = 1, n_embd/n_head
-        q(j,i) = x2((l-1)*n_embd/n_head+j,i)
-    end do
-    end do
-    do i = 1, n_seq
-    do j = 1, n_embd/n_head
-        k(j,i) = kv_cache((l-1)*n_embd/n_head+j,i,1)
-        v(j,i) = kv_cache((l-1)*n_embd/n_head+j,i,2)
-    end do
-    end do
-    call attention(yy, n_embd/n_head, n_seq, n_seq_x, q, k, v, causal_mask)
-    do i = 1, n_seq_x
-    do j = 1, n_embd/n_head
-        y((l-1)*n_embd/n_head+j,i) = yy(j,i)
-    end do
-    end do
-end do
-! Out projection
-y = linear(y, proj_w, proj_b)
-end subroutine
+        x2 = linear(x, attn_w, attn_b)
+        if (use_kv_cache) then
+            do j = 1, n_embd
+                kv_cache(j, n_seq, 1) = x2((2 - 1) * n_embd + j, 1)
+                kv_cache(j, n_seq, 2) = x2((3 - 1) * n_embd + j, 1)
+            end do
+        else
+            do i = 1, n_seq
+                do j = 1, n_embd
+                    kv_cache(j, i, 1) = x2((2 - 1) * n_embd + j, i)
+                    kv_cache(j, i, 2) = x2((3 - 1) * n_embd + j, i)
+                end do
+            end do
+        end if
+        ! Perform attention over each head
+        do l = 1, n_head
+            do i = 1, n_seq_x
+                do j = 1, n_embd / n_head
+                    q(j, i) = x2((l - 1) * n_embd / n_head + j, i)
+                end do
+            end do
+            do i = 1, n_seq
+                do j = 1, n_embd / n_head
+                    k(j, i) = kv_cache((l - 1) * n_embd / n_head + j, i, 1)
+                    v(j, i) = kv_cache((l - 1) * n_embd / n_head + j, i, 2)
+                end do
+            end do
+            call attention(yy, n_embd / n_head, n_seq, n_seq_x, q, k, v, causal_mask)
+            do i = 1, n_seq_x
+                do j = 1, n_embd / n_head
+                    y((l - 1) * n_embd / n_head + j, i) = yy(j, i)
+                end do
+            end do
+        end do
+        ! Out projection
+        y = linear(y, proj_w, proj_b)
+    end subroutine
 
 
-subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, &
-        attn_w, attn_b, attn_proj_w, attn_proj_b, ln1_g, ln1_b, ln2_g, ln2_b, &
-        n_head, use_kv_cache, kv_cache)
-real(sp), intent(inout) :: x(n_embd,n_seq_x)
-real(sp), intent(in) :: &
-    mlp_fc_w(:,:), mlp_fc_b(:), &
-    mlp_proj_w(:,:), mlp_proj_b(:), &
-    attn_w(:,:), attn_b(:), attn_proj_w(:,:), attn_proj_b(:), &
-    ln1_g(:), ln1_b(:), ln2_g(:), ln2_b(:)
-integer, intent(in) :: n_head
-integer, intent(in) :: n_seq, n_seq_x, n_embd
-logical, intent(in) :: use_kv_cache
-real(sp) :: y(n_embd,n_seq_x)
-real(sp) :: yy(n_embd,n_seq_x)
-real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2)
-call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
-call mha(yy, n_seq, n_seq_x, n_embd, y, &
-    attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
-x = x + yy
-!print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b))
-call layer_norm(y, x, ln2_g, ln2_b, 1e-5_sp)
-!print *, "Out1:", y(1,1)
-!x = y
-!print *, "Out2:", x(1,1)
-call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
-end subroutine
+    subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, &
+            attn_w, attn_b, attn_proj_w, attn_proj_b, ln1_g, ln1_b, ln2_g, ln2_b, &
+            n_head, use_kv_cache, kv_cache)
+        real(sp), intent(inout) :: x(n_embd, n_seq_x)
+        real(sp), intent(in) :: &
+                mlp_fc_w(:, :), mlp_fc_b(:), &
+                mlp_proj_w(:, :), mlp_proj_b(:), &
+                attn_w(:, :), attn_b(:), attn_proj_w(:, :), attn_proj_b(:), &
+                ln1_g(:), ln1_b(:), ln2_g(:), ln2_b(:)
+        integer, intent(in) :: n_head
+        integer, intent(in) :: n_seq, n_seq_x, n_embd
+        logical, intent(in) :: use_kv_cache
+        real(sp) :: y(n_embd, n_seq_x)
+        real(sp) :: yy(n_embd, n_seq_x)
+        real(sp), intent(inout) :: kv_cache(n_embd, n_seq, 2)
+        call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp)
+        call mha(yy, n_seq, n_seq_x, n_embd, y, &
+                attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache)
+        x = x + yy
+        !print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b))
+        call layer_norm(y, x, ln2_g, ln2_b, 1e-5_sp)
+        !print *, "Out1:", y(1,1)
+        !x = y
+        !print *, "Out2:", x(1,1)
+        call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b)
+    end subroutine
 
-subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
-        wte, wpe, &
-        mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, &
-        attn_w, attn_b, attn_proj_w, attn_proj_b, &
-        ln1_g, ln1_b, ln2_g, ln2_b, lnf_g, lnf_b, &
-        use_kv_cache, kv_cache)
-integer, intent(in) :: n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head
-integer, intent(in) :: input(n_seq)
-real(sp), intent(in) :: wte(n_embd,n_vocab), wpe(n_embd,n_ctx), &
-    mlp_fc_w(4*n_embd,n_embd,n_layer), mlp_fc_b(4*n_embd,n_layer), &
-    mlp_proj_w(n_embd,4*n_embd,n_layer), mlp_proj_b(n_embd,n_layer), &
-    attn_w(3*n_embd,n_embd,n_layer), attn_b(3*n_embd,n_layer), &
-    attn_proj_w(n_embd,n_embd,n_layer), attn_proj_b(n_embd,n_layer), &
-    ln1_b(n_embd,n_layer), ln1_g(n_embd,n_layer), &
-    ln2_b(n_embd,n_layer), ln2_g(n_embd,n_layer), &
-    lnf_b(n_embd), lnf_g(n_embd)
-logical, intent(in) :: use_kv_cache
-real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2,n_layer)
-real(sp), intent(out) :: y(n_vocab,n_seq_x)
-real(sp) :: x(n_embd,n_seq_x)
-real(sp) :: yy(n_embd,n_seq_x)
-integer :: i, j
-if (use_kv_cache) then
-    i = n_seq
-    do j = 1, n_embd
-        x(j,1) = wte(j,input(i)+1) + wpe(j,i)
-    end do
-else
-    do i = 1, n_seq
-        do j = 1, n_embd
-            x(j,i) = wte(j,input(i)+1) + wpe(j,i)
+    subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, &
+            wte, wpe, &
+            mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, &
+            attn_w, attn_b, attn_proj_w, attn_proj_b, &
+            ln1_g, ln1_b, ln2_g, ln2_b, lnf_g, lnf_b, &
+            use_kv_cache, kv_cache)
+        integer, intent(in) :: n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head
+        integer, intent(in) :: input(n_seq)
+        real(sp), intent(in) :: wte(n_embd, n_vocab), wpe(n_embd, n_ctx), &
+                mlp_fc_w(4 * n_embd, n_embd, n_layer), mlp_fc_b(4 * n_embd, n_layer), &
+                mlp_proj_w(n_embd, 4 * n_embd, n_layer), mlp_proj_b(n_embd, n_layer), &
+                attn_w(3 * n_embd, n_embd, n_layer), attn_b(3 * n_embd, n_layer), &
+                attn_proj_w(n_embd, n_embd, n_layer), attn_proj_b(n_embd, n_layer), &
+                ln1_b(n_embd, n_layer), ln1_g(n_embd, n_layer), &
+                ln2_b(n_embd, n_layer), ln2_g(n_embd, n_layer), &
+                lnf_b(n_embd), lnf_g(n_embd)
+        logical, intent(in) :: use_kv_cache
+        real(sp), intent(inout) :: kv_cache(n_embd, n_seq, 2, n_layer)
+        real(sp), intent(out) :: y(n_vocab, n_seq_x)
+        real(sp) :: x(n_embd, n_seq_x)
+        real(sp) :: yy(n_embd, n_seq_x)
+        integer :: i, j
+        if (use_kv_cache) then
+            i = n_seq
+            do j = 1, n_embd
+                x(j, 1) = wte(j, input(i) + 1) + wpe(j, i)
+            end do
+        else
+            do i = 1, n_seq
+                do j = 1, n_embd
+                    x(j, i) = wte(j, input(i) + 1) + wpe(j, i)
+                end do
+            end do
+        end if
+        !print *, "It fails below:"
+        do j = 1, n_layer
+            i = j
+            !i = 1
+            !    print *, i ! Never gets printed
+            call transformer_block(n_seq, n_seq_x, n_embd, x, &
+                    mlp_fc_w(:, :, i), mlp_fc_b(:, i), &
+                    mlp_proj_w(:, :, i), mlp_proj_b(:, i), &
+                    attn_w(:, :, i), attn_b(:, i), attn_proj_w(:, :, i), attn_proj_b(:, i), &
+                    ln1_g(:, i), ln1_b(:, i), ln2_g(:, i), ln2_b(:, i), &
+                    n_head, use_kv_cache, kv_cache(:, :, :, i))
+            !    print *, x(1,1)
         end do
-    end do
-end if
-!print *, "It fails below:"
-do j = 1, n_layer
-    i = j
-    !i = 1
-!    print *, i ! Never gets printed
-    call transformer_block(n_seq, n_seq_x, n_embd, x, &
-        mlp_fc_w(:,:,i), mlp_fc_b(:,i), &
-        mlp_proj_w(:,:,i), mlp_proj_b(:,i), &
-        attn_w(:,:,i), attn_b(:,i), attn_proj_w(:,:,i), attn_proj_b(:,i), &
-        ln1_g(:,i), ln1_b(:,i), ln2_g(:,i), ln2_b(:,i), &
-        n_head, use_kv_cache, kv_cache(:,:,:,i))
-!    print *, x(1,1)
-end do
-call layer_norm(yy, x, lnf_g, lnf_b, 1e-5)
-x = yy
-!y = matmul(transpose(wte), x)
-call matmul_2d_t(wte, x, y)
-end subroutine
+        call layer_norm(yy, x, lnf_g, lnf_b, 1e-5)
+        x = yy
+        !y = matmul(transpose(wte), x)
+        call matmul_2d_t(wte, x, y)
+    end subroutine
 
-subroutine generate(output, n_tokens_to_generate, m, &
-        n_seq, input, &
-        use_cache, &
-        byte_decoder, stop_text)
-integer, intent(in) :: n_seq, n_tokens_to_generate
-type(model_t), intent(in) :: m
-integer, intent(in) :: input(n_seq)
-logical, intent(in) :: use_cache
-integer, intent(in) :: byte_decoder(:)
-character(*), intent(in), optional :: stop_text ! Stop if you see this text
-integer, allocatable, intent(out) :: output(:)
-real(sp), allocatable :: logits(:,:)
-integer :: i, i1, i2, i3, i4
-integer :: n_seq2, n_seq_x
-integer :: next_id
-integer :: input2(size(input)+n_tokens_to_generate)
-logical :: use_kv_cache
-real(sp) :: kv_cache(m%n_embd,n_seq+n_tokens_to_generate,2,m%n_layer)
-real(sp), allocatable :: kv_cache2(:,:,:,:)
-character(:), allocatable :: output_txt, last_token
-if (present(stop_text)) then
-    output_txt = ""
-end if
-input2(:n_seq) = input
-do i = 1, n_tokens_to_generate
-    if (use_cache) then
-        use_kv_cache = (i > 1) ! Use cache for subsequent tokens
-    else
-        use_kv_cache = .false.
-    end if
-    n_seq2 = n_seq+i-1
-    if (use_kv_cache) then
-        n_seq_x = 1
-    else
-        n_seq_x = n_seq2
-    end if
-    allocate(logits(m%n_vocab, n_seq_x))
-    allocate(kv_cache2(m%n_embd,n_seq2,2,m%n_layer))
-    do i4 = 1, m%n_layer
-    do i3 = 1, 2
-    do i2 = 1, n_seq2
-    do i1 = 1, m%n_embd
-        kv_cache2(i1,i2,i3,i4) = kv_cache(i1,i2,i3,i4)
-    end do
-    end do
-    end do
-    end do
-    call gpt2(logits, m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, &
-            m%n_head, &
-            input2(:n_seq2), &
-            m%wte, m%wpe, &
-            m%mlp_fc_w, m%mlp_fc_b, m%mlp_proj_w, m%mlp_proj_b, &
-            m%attn_w, m%attn_b, m%attn_proj_w, m%attn_proj_b, &
-            m%ln1_g, m%ln1_b, m%ln2_g, m%ln2_b, m%lnf_g, m%lnf_b, use_kv_cache,&
-            kv_cache2)
-    do i4 = 1, m%n_layer
-    do i3 = 1, 2
-    do i2 = 1, n_seq2
-    do i1 = 1, m%n_embd
-        kv_cache(i1,i2,i3,i4) = kv_cache2(i1,i2,i3,i4)
-    end do
-    end do
-    end do
-    end do
-    deallocate(kv_cache2)
-    next_id = maxloc(logits(:,n_seq_x), dim=1)-1
-    input2(n_seq2+1) = next_id
-    last_token = decode([next_id], m%decoder_idx, &
-        m%decoder_txt, byte_decoder)
-    write(*, fmt="(a)", advance="no") last_token
-    if (present(stop_text)) then
-        output_txt = output_txt // last_token
-        if (output_txt(len(output_txt)-len(stop_text)+1:len(output_txt)) == stop_text) then
-            exit
+    subroutine generate(output, n_tokens_to_generate, m, &
+            n_seq, input, &
+            use_cache, &
+            byte_decoder, stop_text)
+        integer, intent(in) :: n_seq, n_tokens_to_generate
+        type(model_t), intent(in) :: m
+        integer, intent(in) :: input(n_seq)
+        logical, intent(in) :: use_cache
+        integer, intent(in) :: byte_decoder(:)
+        character(*), intent(in), optional :: stop_text ! Stop if you see this text
+        integer, allocatable, intent(out) :: output(:)
+        real(sp), allocatable :: logits(:, :)
+        integer :: i, i1, i2, i3, i4
+        integer :: n_seq2, n_seq_x
+        integer :: next_id
+        integer :: input2(size(input) + n_tokens_to_generate)
+        logical :: use_kv_cache
+        real(sp) :: kv_cache(m%n_embd, n_seq + n_tokens_to_generate, 2, m%n_layer)
+        real(sp), allocatable :: kv_cache2(:, :, :, :)
+        character(:), allocatable :: output_txt, last_token
+        if (present(stop_text)) then
+            output_txt = ""
         end if
-    end if
-    deallocate(logits)
-end do
-! output = input2(n_seq+1:n_seq2+1)
-allocate(output(n_seq2 - n_seq + 1))
-do i = 1, n_seq2 - n_seq + 1
-    output(i) = input2(n_seq+i)
-end do
-end subroutine
+        input2(:n_seq) = input
+        do i = 1, n_tokens_to_generate
+            if (use_cache) then
+                use_kv_cache = (i > 1) ! Use cache for subsequent tokens
+            else
+                use_kv_cache = .false.
+            end if
+            n_seq2 = n_seq + i - 1
+            if (use_kv_cache) then
+                n_seq_x = 1
+            else
+                n_seq_x = n_seq2
+            end if
+            allocate(logits(m%n_vocab, n_seq_x))
+            allocate(kv_cache2(m%n_embd, n_seq2, 2, m%n_layer))
+            do i4 = 1, m%n_layer
+                do i3 = 1, 2
+                    do i2 = 1, n_seq2
+                        do i1 = 1, m%n_embd
+                            kv_cache2(i1, i2, i3, i4) = kv_cache(i1, i2, i3, i4)
+                        end do
+                    end do
+                end do
+            end do
+            call gpt2(logits, m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, &
+                    m%n_head, &
+                    input2(:n_seq2), &
+                    m%wte, m%wpe, &
+                    m%mlp_fc_w, m%mlp_fc_b, m%mlp_proj_w, m%mlp_proj_b, &
+                    m%attn_w, m%attn_b, m%attn_proj_w, m%attn_proj_b, &
+                    m%ln1_g, m%ln1_b, m%ln2_g, m%ln2_b, m%lnf_g, m%lnf_b, use_kv_cache, &
+                    kv_cache2)
+            do i4 = 1, m%n_layer
+                do i3 = 1, 2
+                    do i2 = 1, n_seq2
+                        do i1 = 1, m%n_embd
+                            kv_cache(i1, i2, i3, i4) = kv_cache2(i1, i2, i3, i4)
+                        end do
+                    end do
+                end do
+            end do
+            deallocate(kv_cache2)
+            next_id = maxloc(logits(:, n_seq_x), dim = 1) - 1
+            input2(n_seq2 + 1) = next_id
+            last_token = decode([next_id], m%decoder_idx, &
+                    m%decoder_txt, byte_decoder)
+            write(*, fmt = "(a)", advance = "no") last_token
+            if (present(stop_text)) then
+                output_txt = output_txt // last_token
+                if (output_txt(len(output_txt) - len(stop_text) + 1:len(output_txt)) == stop_text) then
+                    exit
+                end if
+            end if
+            deallocate(logits)
+        end do
+        ! output = input2(n_seq+1:n_seq2+1)
+        allocate(output(n_seq2 - n_seq + 1))
+        do i = 1, n_seq2 - n_seq + 1
+            output(i) = input2(n_seq + i)
+        end do
+    end subroutine
 
 end module
diff --git a/lp-fastgpt/README.md b/lp-fastgpt/README.md
new file mode 100644
index 0000000..e69de29
diff --git a/lp-fastgpt/create_model.py b/lp-fastgpt/create_model.py
new file mode 100644
index 0000000..d07a740
--- /dev/null
+++ b/lp-fastgpt/create_model.py
@@ -0,0 +1,355 @@
+"""This script loads the specified GPT-2 model from OpenAI
+using TensorFlow, converts it into our custom format and saves
+it to `model.dat`, which contains everything (all the
+parameters, all the weights, encoding/decoding information).
+
+Parts of this script were taken from the picoGPT project:
+https://github.com/jaymody/picoGPT
+
+Those are licensed as:
+
+MIT License
+
+Copyright (c) 2023 Jay Mody
+
+Permission is hereby granted, free of charge, to any person
+obtaining a copy of this software and associated documentation
+files (the "Software"), to deal in the Software without
+restriction, including without limitation the rights to use,
+copy, modify, merge, publish, distribute, sublicense, and/or
+sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the
+Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
+KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
+PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+"""
+
+from time import monotonic as clock
+import os
+import json
+import re
+
+
+import numpy as np
+import requests
+import tensorflow as tf
+from tqdm import tqdm
+
+
+from typing import Optional
+
+
+def download_gpt2_files(model_size, model_dir):
+    assert model_size in ["124M", "355M", "774M", "1558M"]
+    for filename in [
+        "checkpoint",
+        "encoder.json",
+        "hparams.json",
+        "model.ckpt.data-00000-of-00001",
+        "model.ckpt.index",
+        "model.ckpt.meta",
+        "vocab.bpe",
+    ]:
+        url = "https://openaipublic.blob.core.windows.net/gpt-2/models"
+        r = requests.get(f"{url}/{model_size}/{filename}", stream=True)
+        r.raise_for_status()
+
+        with open(os.path.join(model_dir, filename), "wb") as f:
+            file_size = int(r.headers["content-length"])
+            chunk_size = 1000
+            with tqdm(
+                    ncols=100,
+                    desc="Fetching " + filename,
+                    total=file_size,
+                    unit_scale=True,
+            ) as pbar:
+                # 1k for chunk_size, since Ethernet packet size is around 1500 bytes
+                for chunk in r.iter_content(chunk_size=chunk_size):
+                    f.write(chunk)
+                    pbar.update(chunk_size)
+
+
+def load_gpt2_params_from_tf_ckpt(tf_ckpt_path, hparams):
+    def set_in_nested_dict(d, keys, val):
+        if not keys:
+            return val
+        if keys[0] not in d:
+            d[keys[0]] = {}
+        d[keys[0]] = set_in_nested_dict(d[keys[0]], keys[1:], val)
+        return d
+
+    init_vars = tf.train.list_variables(tf_ckpt_path)
+    params = {"blocks": [{} for _ in range(hparams["n_layer"])]}
+    for name, _ in init_vars:
+        array = np.squeeze(tf.train.load_variable(tf_ckpt_path, name))
+        name = name.removeprefix("model/")
+        if name.startswith("h"):
+            m = re.match(r"h([0-9]+)/(.*)", name)
+            n = int(m[1])
+            sub_name = m[2]
+            set_in_nested_dict(params["blocks"][n], sub_name.split("/"), array)
+        else:
+            set_in_nested_dict(params, name.split("/"), array)
+
+    return params
+
+
+def tf_load_encoder_hparams_and_params(model_size, models_dir):
+
+    assert model_size in ["124M", "355M", "774M", "1558M"]
+
+    model_dir = os.path.join(models_dir, model_size)
+
+    tf_ckpt_path = tf.train.latest_checkpoint(model_dir)
+
+    if not tf_ckpt_path:  # download files if necessary
+        os.makedirs(model_dir, exist_ok=True)
+        download_gpt2_files(model_size, model_dir)
+        tf_ckpt_path = tf.train.latest_checkpoint(model_dir)
+
+    hparams = json.load(open(os.path.join(model_dir, "hparams.json")))
+
+    params = load_gpt2_params_from_tf_ckpt(tf_ckpt_path, hparams)
+
+    return hparams, params
+
+
+from dataclasses import dataclass
+
+
+@dataclass
+class Model:
+    blocks      : dict[str, dict[str, dict[str, np.ndarray]]]
+    n_embd      : int
+    n_layer     : int
+    mlp_fc_w    : np.ndarray
+    mlp_fc_b    : np.ndarray
+    mlp_proj_w  : np.ndarray
+    mlp_proj_b  : np.ndarray
+    attn_w      : np.ndarray
+    attn_b      : np.ndarray
+    attn_proj_w : np.ndarray
+    attn_proj_b : np.ndarray
+    ln1_g       : np.ndarray
+    ln1_b       : np.ndarray
+    ln2_g       : np.ndarray
+    ln2_b       : np.ndarray
+    wte         : np.ndarray
+    wpe         : np.ndarray
+    lnf_g       : np.ndarray
+    lnf_b       : np.ndarray
+
+
+def convert(params,
+            n_head,
+            n_ctx,
+            idx,
+            decoder_txt,
+            vocab_idx,
+            vocab_txt,
+            byte_decoder) -> Model:
+
+    t1 = clock()
+
+    # must predefine just to get shapes ...
+    blocks  = params["blocks"]
+    n_embd  = blocks[0]["ln_1"]["b"].size
+    n_layer = len(blocks)
+
+    m : Model = Model(
+        blocks      = blocks,
+        n_embd      = n_embd,
+        n_layer     = n_layer,
+
+        mlp_fc_w    = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32),
+        mlp_fc_b    = np.empty((n_layer, 4 * n_embd)         , dtype=np.float32),
+        mlp_proj_w  = np.empty((n_layer, 4 * n_embd, n_embd) , dtype=np.float32),
+        mlp_proj_b  = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        attn_w      = np.empty((n_layer, n_embd, 3 * n_embd) , dtype=np.float32),
+        attn_b      = np.empty((n_layer, 3 * n_embd)         , dtype=np.float32),
+        attn_proj_w = np.empty((n_layer, n_embd, n_embd)     , dtype=np.float32),
+        attn_proj_b = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln1_g       = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln1_b       = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln2_g       = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln2_b       = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        wte         = np.empty(0),
+        wpe         = np.empty(0),
+        lnf_g       = np.empty(0),
+        lnf_b       = np.empty(0),
+    )
+    for i, block in enumerate(blocks):
+        m.mlp_fc_w[i, :, :]    = block["mlp"]["c_fc"]["w"]
+        m.mlp_fc_b[i, :]       = block["mlp"]["c_fc"]["b"]
+        m.mlp_proj_w[i, :, :]  = block["mlp"]["c_proj"]["w"]
+        m.mlp_proj_b[i, :]     = block["mlp"]["c_proj"]["b"]
+        m.attn_w[i, :, :]      = block["attn"]["c_attn"]["w"]
+        m.attn_b[i, :]         = block["attn"]["c_attn"]["b"]
+        m.attn_proj_w[i, :, :] = block["attn"]["c_proj"]["w"]
+        m.attn_proj_b[i, :]    = block["attn"]["c_proj"]["b"]
+        m.ln1_g[i, :]          = block["ln_1"]["g"]
+        m.ln1_b[i, :]          = block["ln_1"]["b"]
+        m.ln2_g[i, :]          = block["ln_2"]["g"]
+        m.ln2_b[i, :]          = block["ln_2"]["b"]
+    m.wte = params["wte"]
+    m.wpe = params["wpe"]
+    m.lnf_g = params["ln_f"]["g"]
+    m.lnf_b = params["ln_f"]["b"]
+    t2 = clock()
+    print("Transform time: ", t2 - t1)
+    t1 = clock()
+
+    n_vocab = np.size(m.wte, 0)
+    assert np.size(m.wte, 1) == n_embd
+
+    model_type = 0xfa51697  # fastGPT
+    model_version = 1
+
+    # Save the model
+    f = open("model.dat", "w")
+
+    # what is this? just gets thrown away, no ? ...
+
+    np.array([model_type,
+              model_version,
+              n_vocab,
+              n_ctx,
+              n_embd,
+              n_layer,
+              n_head,
+              len(idx),
+              len(decoder_txt.encode("utf-8")),
+              len(vocab_idx),
+              len(vocab_txt.encode("utf-8")),
+              len(byte_decoder)],
+             dtype=np.int32).tofile(f)
+
+    m.wte.tofile(f)
+    m.wpe.tofile(f)
+    m.mlp_fc_w.tofile(f)
+    m.mlp_fc_b.tofile(f)
+    m.mlp_proj_w.tofile(f)
+
+    m.mlp_proj_b.tofile(f)
+    m.attn_w.tofile(f)
+    m.attn_b.tofile(f)
+    m.attn_proj_w.tofile(f)
+    m.attn_proj_b.tofile(f)
+    m.ln1_b.tofile(f)
+    m.ln1_g.tofile(f)
+    m.ln2_b.tofile(f)
+    m.ln2_g.tofile(f)
+    m.lnf_b.tofile(f)
+    m.lnf_g.tofile(f)
+
+    idx.tofile(f)
+
+    f.write(decoder_txt)
+
+    vocab_idx.tofile(f)
+
+    f.write(vocab_txt)
+
+    byte_decoder.tofile(f)
+
+    t2 = clock()
+    print("Save time: ", t2 - t1)
+
+    return m
+
+
+def load_decoder(filename):
+    D = json.load(open(filename))
+    D2 = {v: k for k, v in D.items()}
+    i = 0
+    decoder = []
+    while True:
+        if i not in D2:
+            break
+        decoder.append(D2[i])
+        i += 1
+    return decoder
+
+
+def load_vocab(filename):
+    D = open(filename).read()
+    D = D.split("\n")
+    D = D[1:]
+    return D
+
+
+def decoder_idx(decoder):
+    i = 0
+    idx = np.empty(len(decoder) + 1, dtype=np.int32)
+    idx[0] = i
+    for n, t in enumerate(decoder):
+        i += len(t.encode("utf-8"))
+        idx[n + 1] = i
+    assert idx[-1] == len("".join(decoder).encode("utf-8"))
+    return idx
+
+
+def bytes_to_unicode():
+    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    cs = bs[:]
+    n = 0
+    for b in range(2 ** 8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2 ** 8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    btu = dict(zip(bs, cs))
+    byte_decoder = {v: k for k, v in btu.items()}
+    bd = np.zeros(324, dtype=np.int32)
+    for y in byte_decoder:
+        x = ord(y)
+        bd[x] = byte_decoder[y]
+    bd2 = np.zeros(256, dtype=np.int32)
+    for i in range(np.size(bd)):
+        bd2[bd[i]] = i
+    return bd2
+
+
+def main(model_size: str = "124M", models_dir: str = "models") -> Model:
+    print("Loading model")
+    # load encoder, hparams, and params from the released open-ai gpt-2 files
+    t1 = clock()
+    hparams, params = tf_load_encoder_hparams_and_params(model_size, models_dir)
+    decoder = load_decoder(os.path.join(models_dir, model_size, "encoder.json"))
+    vocab   = load_vocab(os.path.join(models_dir, model_size, "vocab.bpe"))
+    t2 = clock()
+    print("  Done. Loading time: ", t2 - t1)
+
+    # generate output ids
+    print("Converting model, saving to `model.dat`")
+    t1 = clock()
+    decoder_txt = "".join(decoder)
+    idx = decoder_idx(decoder)
+    vocab_txt = "".join(vocab)
+    vocab_idx = decoder_idx(vocab)
+    byte_decoder = bytes_to_unicode()
+    m : Model = \
+        convert(params, hparams["n_head"], hparams["n_ctx"], idx, decoder_txt,
+            vocab_idx, vocab_txt, byte_decoder)
+    t2 = clock()
+    print("  Done. Time: ", t2 - t1)
+    return m
+
+
+if __name__ == "__main__":
+    import fire
+
+    fire.Fire(main)
diff --git a/lp-fastgpt/main.py b/lp-fastgpt/main.py
new file mode 100644
index 0000000..71245a4
--- /dev/null
+++ b/lp-fastgpt/main.py
@@ -0,0 +1,21 @@
+import create_model
+
+import numpy as np
+
+
+def load_input(filename: str, ) -> tuple[str, int]:
+    """Straight transcription of load_input from
+    driver.f90."""
+    n_tokens_to_generate: int = 20
+    input_txt: str = '''Alan Turing theorized that
+    computers would one day become very powerful
+    , but even he could not imagine'''
+    # input_txt2: str = ''
+    # u : int = 0
+    # ios : int = 0
+    return input_txt, n_tokens_to_generate
+
+
+if __name__ == '__main__':
+    print("hello lp_fastgpt")
+    create_model.main()
diff --git a/lp-fastgpt/requirements.txt b/lp-fastgpt/requirements.txt
new file mode 100644
index 0000000..e69de29
diff --git a/main.f90 b/main.f90
index dcb8ff1..d551d2f 100644
--- a/main.f90
+++ b/main.f90
@@ -1,7 +1,7 @@
 program gpt2
-use driver, only: gpt2_driver, model_t
-implicit none
-integer, allocatable :: input(:), output(:)
-type(model_t) :: m
-call gpt2_driver(input, output, m)
+    use driver, only : gpt2_driver, model_t
+    implicit none
+    integer, allocatable :: input(:), output(:)
+    type(model_t) :: m
+    call gpt2_driver(input, output, m)
 end program

From f6ccd3768b9cbfcaabf166e3c85802cfb73a30a2 Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Thu, 3 Aug 2023 18:07:46 -0700
Subject: [PATCH 74/84] many asserts -- round-tripping is next

---
 lp-fastgpt/create_model.py | 330 ++++++++++++++++++++++++++-----------
 1 file changed, 237 insertions(+), 93 deletions(-)

diff --git a/lp-fastgpt/create_model.py b/lp-fastgpt/create_model.py
index d07a740..8653281 100644
--- a/lp-fastgpt/create_model.py
+++ b/lp-fastgpt/create_model.py
@@ -48,7 +48,7 @@
 from tqdm import tqdm
 
 
-from typing import Optional
+from typing import Optional, Union, Any
 
 
 def download_gpt2_files(model_size, model_dir):
@@ -128,12 +128,67 @@ def tf_load_encoder_hparams_and_params(model_size, models_dir):
 
 from dataclasses import dataclass
 
+HParamsType = dict[str, int]
+
+# === Block-Attention components ======================
+CAttnBType         = np.ndarray
+CAttnBShape        = (     2304,)
+CAttnWType         = np.ndarray
+CAttnWShape        = (768, 2304,)
+CAttnType          = dict[str, Union[CAttnBType, CAttnWType]]
+
+CProjBType         = np.ndarray
+CProjBShape        = (      768,)
+CProjWType         = np.ndarray
+CProjWShape        = (768,  768,)
+CProjType          = dict[str, Union[CProjBType, CProjWType]]
+# --- top level
+BlockAttnType      = dict[str, Union[CAttnType, CProjType]]
+# === BlockIn components ==============================
+BlockLnBType       = np.ndarray
+BlockLnBShape      = (768,)
+BlockLnGType       = np.ndarray
+BlockLnGShape      = (768,)
+# --- top level
+BlockLnType        = dict[str, Union[BlockLnBType, BlockLnGType]]
+# === BlockMlp components =============================
+MlpCFcBType        = np.ndarray
+MlpCFcBShape       = (     3072,)
+MlpCFcWType        = np.ndarray
+MlpCFcWShape       = (768, 3072,)
+MlpCFcType         = dict[str, Union[MlpCFcBType, MlpCFcWType]]
+
+MlpCProjBType      = np.ndarray
+MlpCProjBShape     = (      768,)
+MlpCProjWType      = np.ndarray
+MlpCProjWShape     = (3072, 768,)
+MlpCProjType       = dict[str, Union[MlpCProjBType, MlpCProjWType]]
+# --- top level
+BlockMlpType       = dict[str, Union[MlpCFcType, MlpCProjType]]
+# =====================================================
+ParamsBlockType    = dict[str, Union[BlockAttnType,
+                                     BlockLnType,  # two of these
+                                     BlockMlpType]]
+#                 len=12
+ParamsBlocksType   = list[ParamsBlockType]
+ParamsLnFType      = dict[str, np.ndarray]
+ParamsLnFValShape  = BlockLnBShape
+ParamsWpeType      = np.ndarray
+ParamsWpeShape     = ( 1024, 768,)
+ParamsWteType      = np.ndarray
+ParamsWteShape     = (50257, 768,)
+ParamsType         = dict[str, Union[ParamsBlocksType,
+                                     ParamsLnFType,
+                                     ParamsWpeType,
+                                     ParamsWteType]]
+
 
 @dataclass
 class Model:
-    blocks      : dict[str, dict[str, dict[str, np.ndarray]]]
+    blocks      : ParamsBlocksType  # TODO: take this out of the model
     n_embd      : int
     n_layer     : int
+    # check shapes in asserts for now; type system too weak.
     mlp_fc_w    : np.ndarray
     mlp_fc_b    : np.ndarray
     mlp_proj_w  : np.ndarray
@@ -164,110 +219,161 @@ def convert(params,
     t1 = clock()
 
     # must predefine just to get shapes ...
-    blocks  = params["blocks"]
+    blocks : ParamsBlocksType = params["blocks"]
+    nblocks = len(blocks)
+    assert nblocks == 12
+
     n_embd  = blocks[0]["ln_1"]["b"].size
-    n_layer = len(blocks)
+    assert n_embd == 768
 
-    m : Model = Model(
-        blocks      = blocks,
-        n_embd      = n_embd,
-        n_layer     = n_layer,
+    n_layer = nblocks
+    assert n_layer == 12
+
+    mo : Model = make_empty_model(blocks, n_embd, n_layer)
 
-        mlp_fc_w    = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32),
-        mlp_fc_b    = np.empty((n_layer, 4 * n_embd)         , dtype=np.float32),
-        mlp_proj_w  = np.empty((n_layer, 4 * n_embd, n_embd) , dtype=np.float32),
-        mlp_proj_b  = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        attn_w      = np.empty((n_layer, n_embd, 3 * n_embd) , dtype=np.float32),
-        attn_b      = np.empty((n_layer, 3 * n_embd)         , dtype=np.float32),
-        attn_proj_w = np.empty((n_layer, n_embd, n_embd)     , dtype=np.float32),
-        attn_proj_b = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        ln1_g       = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        ln1_b       = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        ln2_g       = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        ln2_b       = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        wte         = np.empty(0),
-        wpe         = np.empty(0),
-        lnf_g       = np.empty(0),
-        lnf_b       = np.empty(0),
-    )
     for i, block in enumerate(blocks):
-        m.mlp_fc_w[i, :, :]    = block["mlp"]["c_fc"]["w"]
-        m.mlp_fc_b[i, :]       = block["mlp"]["c_fc"]["b"]
-        m.mlp_proj_w[i, :, :]  = block["mlp"]["c_proj"]["w"]
-        m.mlp_proj_b[i, :]     = block["mlp"]["c_proj"]["b"]
-        m.attn_w[i, :, :]      = block["attn"]["c_attn"]["w"]
-        m.attn_b[i, :]         = block["attn"]["c_attn"]["b"]
-        m.attn_proj_w[i, :, :] = block["attn"]["c_proj"]["w"]
-        m.attn_proj_b[i, :]    = block["attn"]["c_proj"]["b"]
-        m.ln1_g[i, :]          = block["ln_1"]["g"]
-        m.ln1_b[i, :]          = block["ln_1"]["b"]
-        m.ln2_g[i, :]          = block["ln_2"]["g"]
-        m.ln2_b[i, :]          = block["ln_2"]["b"]
-    m.wte = params["wte"]
-    m.wpe = params["wpe"]
-    m.lnf_g = params["ln_f"]["g"]
-    m.lnf_b = params["ln_f"]["b"]
+        mo.mlp_fc_w[i, :, :]    = block["mlp"]["c_fc"]["w"]
+        mo.mlp_fc_b[i, :]       = block["mlp"]["c_fc"]["b"]
+        mo.mlp_proj_w[i, :, :]  = block["mlp"]["c_proj"]["w"]
+        mo.mlp_proj_b[i, :]     = block["mlp"]["c_proj"]["b"]
+        mo.attn_w[i, :, :]      = block["attn"]["c_attn"]["w"]
+        mo.attn_b[i, :]         = block["attn"]["c_attn"]["b"]
+        mo.attn_proj_w[i, :, :] = block["attn"]["c_proj"]["w"]
+        mo.attn_proj_b[i, :]    = block["attn"]["c_proj"]["b"]
+
+        mo.ln1_g[i, :]          = block["ln_1"]["g"]
+        mo.ln1_b[i, :]          = block["ln_1"]["b"]
+        mo.ln2_g[i, :]          = block["ln_2"]["g"]
+        mo.ln2_b[i, :]          = block["ln_2"]["b"]
+
+    mo.wte   = params["wte"]
+    mo.wpe   = params["wpe"]
+    mo.lnf_g = params["ln_f"]["g"]
+    mo.lnf_b = params["ln_f"]["b"]
+
     t2 = clock()
     print("Transform time: ", t2 - t1)
-    t1 = clock()
 
-    n_vocab = np.size(m.wte, 0)
-    assert np.size(m.wte, 1) == n_embd
+    t1 = clock()
 
+    n_vocab = np.size(mo.wte, 0)
     model_type = 0xfa51697  # fastGPT
     model_version = 1
 
     # Save the model
     f = open("model.dat", "w")
 
-    # what is this? just gets thrown away, no ? ...
-
-    np.array([model_type,
-              model_version,
-              n_vocab,
-              n_ctx,
-              n_embd,
-              n_layer,
-              n_head,
-              len(idx),
-              len(decoder_txt.encode("utf-8")),
-              len(vocab_idx),
-              len(vocab_txt.encode("utf-8")),
-              len(byte_decoder)],
-             dtype=np.int32).tofile(f)
-
-    m.wte.tofile(f)
-    m.wpe.tofile(f)
-    m.mlp_fc_w.tofile(f)
-    m.mlp_fc_b.tofile(f)
-    m.mlp_proj_w.tofile(f)
-
-    m.mlp_proj_b.tofile(f)
-    m.attn_w.tofile(f)
-    m.attn_b.tofile(f)
-    m.attn_proj_w.tofile(f)
-    m.attn_proj_b.tofile(f)
-    m.ln1_b.tofile(f)
-    m.ln1_g.tofile(f)
-    m.ln2_b.tofile(f)
-    m.ln2_g.tofile(f)
-    m.lnf_b.tofile(f)
-    m.lnf_g.tofile(f)
 
-    idx.tofile(f)
 
-    f.write(decoder_txt)
+    model_metadata = np.array(
+        [model_type,
+         model_version,
+         n_vocab,
+         n_ctx,
+         n_embd,
+         n_layer,
+         n_head,
+         len(idx),
+         len(decoder_txt.encode("utf-8")),
+         len(vocab_idx),
+         len(vocab_txt.encode("utf-8")),
+         len(byte_decoder)], dtype=np.int32)
+
+    model_metadata.tofile(f)
+
+    mo.wte.tofile(f)
+    mo.wpe.tofile(f)
+
+    mo.mlp_fc_w.tofile(f)
+    mo.mlp_fc_b.tofile(f)
+
+    mo.mlp_proj_w.tofile(f)
+    mo.mlp_proj_b.tofile(f)
+
+    mo.attn_w.tofile(f)
+    mo.attn_b.tofile(f)
+
+    mo.attn_proj_w.tofile(f)
+    mo.attn_proj_b.tofile(f)
+
+    mo.ln1_b.tofile(f)
+    mo.ln1_g.tofile(f)
+    mo.ln2_b.tofile(f)
+    mo.ln2_g.tofile(f)
+    mo.lnf_b.tofile(f)
+    mo.lnf_g.tofile(f)
 
+    idx.tofile(f)
+    f.write(decoder_txt)
     vocab_idx.tofile(f)
 
     f.write(vocab_txt)
 
     byte_decoder.tofile(f)
 
+    check_model_shapes(
+        mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks)
+
     t2 = clock()
     print("Save time: ", t2 - t1)
 
-    return m
+    t1 = clock()
+    m = make_empty_model(blocks, n_embd, n_layer)
+    t2 = clock()
+    print("Restore time: ", t2 - t1)
+
+
+    return mo
+
+
+def make_empty_model(blocks, n_embd, n_layer):
+    mo: Model = Model(
+        blocks  = blocks,
+        n_embd  = n_embd,
+        n_layer = n_layer,
+
+        mlp_fc_w    = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32),
+        mlp_fc_b    = np.empty((n_layer, 4 * n_embd)         , dtype=np.float32),
+        mlp_proj_w  = np.empty((n_layer, 4 * n_embd, n_embd) , dtype=np.float32),
+        mlp_proj_b  = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        attn_w      = np.empty((n_layer, n_embd, 3 * n_embd) , dtype=np.float32),
+        attn_b      = np.empty((n_layer, 3 * n_embd)         , dtype=np.float32),
+        attn_proj_w = np.empty((n_layer, n_embd, n_embd)     , dtype=np.float32),
+        attn_proj_b = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln1_g       = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln1_b       = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln2_g       = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln2_b       = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        wte         = np.empty(0                             , dtype=np.float32),
+        wpe         = np.empty(0                             , dtype=np.float32),
+        lnf_g       = np.empty(0                             , dtype=np.float32),
+        lnf_b       = np.empty(0                             , dtype=np.float32),
+    )
+    return mo
+
+
+def check_model_shapes(mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks):
+    assert mo.mlp_fc_w.shape    == (nblocks,) + MlpCFcWShape
+    assert mo.mlp_fc_b.shape    == (nblocks,) + MlpCFcBShape
+    assert mo.mlp_proj_w.shape  == (nblocks,) + MlpCProjWShape
+    assert mo.mlp_proj_b.shape  == (nblocks,) + MlpCProjBShape
+    assert mo.attn_w.shape      == (nblocks,) + CAttnWShape
+    assert mo.attn_b.shape      == (nblocks,) + CAttnBShape
+    assert mo.attn_proj_w.shape == (nblocks,) + CProjWShape
+    assert mo.attn_proj_b.shape == (nblocks,) + CProjBShape
+    assert mo.ln1_g.shape       == (nblocks,) + BlockLnGShape
+    assert mo.ln1_b.shape       == (nblocks,) + BlockLnBShape
+    assert mo.ln2_g.shape       == (nblocks,) + BlockLnGShape
+    assert mo.ln2_b.shape       == (nblocks,) + BlockLnBShape
+    assert mo.wte.shape         == ParamsWteShape
+    assert mo.wpe.shape         == ParamsWpeShape
+    assert mo.lnf_g.shape       == BlockLnGShape
+    assert mo.lnf_b.shape       == BlockLnBShape
+    assert n_vocab              == ParamsWteShape[0] == 50_257
+    assert np.size(mo.wte, 1)   == n_embd == 768
+    assert idx.shape            == (50_258,)
+    assert vocab_idx.shape      == (50_002,)
+    assert byte_decoder.shape   == (256,)
 
 
 def load_decoder(filename):
@@ -302,7 +408,9 @@ def decoder_idx(decoder):
 
 
 def bytes_to_unicode():
-    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    bs = list(range(ord("!"), ord("~") + 1)) + \
+        list(range(ord("¡"), ord("¬") + 1)) + \
+        list(range(ord("®"), ord("ÿ") + 1))
     cs = bs[:]
     n = 0
     for b in range(2 ** 8):
@@ -323,29 +431,65 @@ def bytes_to_unicode():
     return bd2
 
 
-def main(model_size: str = "124M", models_dir: str = "models") -> Model:
-    print("Loading model")
+def main(model_size: str = "124M",
+         models_dir: str = "models") -> Model:
+
+    # ================================================================
     # load encoder, hparams, and params from the released open-ai gpt-2 files
+    print("Loading model")
     t1 = clock()
-    hparams, params = tf_load_encoder_hparams_and_params(model_size, models_dir)
-    decoder = load_decoder(os.path.join(models_dir, model_size, "encoder.json"))
-    vocab   = load_vocab(os.path.join(models_dir, model_size, "vocab.bpe"))
+
+    hparams : HParamsType
+    params  : ParamsType
+
+    hparams, params = \
+        tf_load_encoder_hparams_and_params(model_size, models_dir)
+
+    decoder : list[str] = \
+        load_decoder(os.path.join(models_dir, model_size, "encoder.json"))
+
+    assert len(decoder) == 50_257
+
+    vocab : list[str] = \
+        load_vocab(os.path.join(models_dir, model_size, "vocab.bpe"))
+
+    assert len(vocab) == 50_001
+
     t2 = clock()
     print("  Done. Loading time: ", t2 - t1)
-
+    # ================================================================
     # generate output ids
     print("Converting model, saving to `model.dat`")
     t1 = clock()
-    decoder_txt = "".join(decoder)
-    idx = decoder_idx(decoder)
-    vocab_txt = "".join(vocab)
-    vocab_idx = decoder_idx(vocab)
+
+    decoder_txt : str = "".join(decoder)
+
+    idx : np.ndarray = decoder_idx(decoder)
+
+    vocab_txt    = "".join(vocab)
+
+    vocab_idx    = decoder_idx(vocab)
+
+    assert vocab_idx.shape == (50_002,)
+
     byte_decoder = bytes_to_unicode()
+
+    assert byte_decoder.shape == (256,)
+
     m : Model = \
-        convert(params, hparams["n_head"], hparams["n_ctx"], idx, decoder_txt,
-            vocab_idx, vocab_txt, byte_decoder)
+        convert(params,
+                hparams["n_head"],
+                hparams["n_ctx"],
+                idx,
+                decoder_txt,
+                vocab_idx,
+                vocab_txt,
+                byte_decoder)
+
     t2 = clock()
     print("  Done. Time: ", t2 - t1)
+    # ================================================================
+
     return m
 
 

From 2797bb23c558801049781557843e701fb0b22c5c Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Thu, 3 Aug 2023 18:46:31 -0700
Subject: [PATCH 75/84] all asserts -- factoring metadata is next

---
 lp-fastgpt/create_model.py | 197 +++++++++++++++++++++++++------------
 1 file changed, 136 insertions(+), 61 deletions(-)

diff --git a/lp-fastgpt/create_model.py b/lp-fastgpt/create_model.py
index 8653281..c632e60 100644
--- a/lp-fastgpt/create_model.py
+++ b/lp-fastgpt/create_model.py
@@ -186,8 +186,19 @@ def tf_load_encoder_hparams_and_params(model_size, models_dir):
 @dataclass
 class Model:
     blocks      : ParamsBlocksType  # TODO: take this out of the model
-    n_embd      : int
-    n_layer     : int
+    # integer metadata
+    model_type       : int
+    model_version    : int
+    n_vocab          : int
+    n_ctx            : int
+    n_embd           : int
+    n_layer          : int
+    n_head           : int
+    idx_len          : int
+    decoder_txt_len  : int
+    vocab_idx_len    : int
+    vocab_txt_len    : int
+    byte_decoder_len : int
     # check shapes in asserts for now; type system too weak.
     mlp_fc_w    : np.ndarray
     mlp_fc_b    : np.ndarray
@@ -224,12 +235,28 @@ def convert(params,
     assert nblocks == 12
 
     n_embd  = blocks[0]["ln_1"]["b"].size
-    assert n_embd == 768
-
     n_layer = nblocks
     assert n_layer == 12
 
-    mo : Model = make_empty_model(blocks, n_embd, n_layer)
+    n_vocab = ParamsWteShape[0]  # np.size(mo.wte, 0)
+    model_type = 0xfa51697  # fastGPT
+    model_version = 1
+
+    mo : Model = make_empty_model_with_metadata(
+        blocks           = blocks,
+        model_type       = model_type,
+        model_version    = model_version,
+        n_vocab          = n_vocab,
+        n_ctx            = n_ctx,
+        n_embd           = n_embd,
+        n_layer          = n_layer,
+        n_head           = n_head,
+        idx_len          = len(idx),
+        decoder_txt_len  = len(decoder_txt.encode("utf-8")),
+        vocab_idx_len    = len(vocab_idx),
+        vocab_txt_len    = len(vocab_txt.encode("utf-8")),
+        byte_decoder_len = len(byte_decoder),
+    )
 
     for i, block in enumerate(blocks):
         mo.mlp_fc_w[i, :, :]    = block["mlp"]["c_fc"]["w"]
@@ -254,71 +281,82 @@ def convert(params,
     t2 = clock()
     print("Transform time: ", t2 - t1)
 
-    t1 = clock()
-
-    n_vocab = np.size(mo.wte, 0)
-    model_type = 0xfa51697  # fastGPT
-    model_version = 1
+    check_model(
+        mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks)
 
     # Save the model
-    f = open("model.dat", "w")
-
-
+    t1 = clock()
+    with open("model.dat", "w") as f:
 
-    model_metadata = np.array(
-        [model_type,
-         model_version,
-         n_vocab,
-         n_ctx,
-         n_embd,
-         n_layer,
-         n_head,
-         len(idx),
-         len(decoder_txt.encode("utf-8")),
-         len(vocab_idx),
-         len(vocab_txt.encode("utf-8")),
-         len(byte_decoder)], dtype=np.int32)
+        model_metadata = np.array(
+            [model_type,
+             model_version,
+             n_vocab,
+             n_ctx,
+             n_embd,
+             n_layer,
+             n_head,
+             len(idx),
+             len(decoder_txt.encode("utf-8")),
+             len(vocab_idx),
+             len(vocab_txt.encode("utf-8")),
+             len(byte_decoder)], dtype=np.int32)
 
-    model_metadata.tofile(f)
+        model_metadata.tofile(f)
 
-    mo.wte.tofile(f)
-    mo.wpe.tofile(f)
+        mo.wte.tofile(f)
+        mo.wpe.tofile(f)
 
-    mo.mlp_fc_w.tofile(f)
-    mo.mlp_fc_b.tofile(f)
+        mo.mlp_fc_w.tofile(f)
+        mo.mlp_fc_b.tofile(f)
 
-    mo.mlp_proj_w.tofile(f)
-    mo.mlp_proj_b.tofile(f)
+        mo.mlp_proj_w.tofile(f)
+        mo.mlp_proj_b.tofile(f)
 
-    mo.attn_w.tofile(f)
-    mo.attn_b.tofile(f)
+        mo.attn_w.tofile(f)
+        mo.attn_b.tofile(f)
 
-    mo.attn_proj_w.tofile(f)
-    mo.attn_proj_b.tofile(f)
+        mo.attn_proj_w.tofile(f)
+        mo.attn_proj_b.tofile(f)
 
-    mo.ln1_b.tofile(f)
-    mo.ln1_g.tofile(f)
-    mo.ln2_b.tofile(f)
-    mo.ln2_g.tofile(f)
-    mo.lnf_b.tofile(f)
-    mo.lnf_g.tofile(f)
+        mo.ln1_b.tofile(f)
+        mo.ln1_g.tofile(f)
+        mo.ln2_b.tofile(f)
+        mo.ln2_g.tofile(f)
+        mo.lnf_b.tofile(f)
+        mo.lnf_g.tofile(f)
 
-    idx.tofile(f)
-    f.write(decoder_txt)
-    vocab_idx.tofile(f)
+        idx.tofile(f)
+        f.write(decoder_txt)
 
-    f.write(vocab_txt)
+        vocab_idx.tofile(f)
 
-    byte_decoder.tofile(f)
+        f.write(vocab_txt)
 
-    check_model_shapes(
-        mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks)
+        byte_decoder.tofile(f)
 
     t2 = clock()
     print("Save time: ", t2 - t1)
 
     t1 = clock()
-    m = make_empty_model(blocks, n_embd, n_layer)
+    m = make_empty_model_with_metadata(
+        blocks           = blocks,
+        model_type       = model_type,
+        model_version    = model_version,
+        n_vocab          = n_vocab,
+        n_ctx            = n_ctx,
+        n_embd           = n_embd,
+        n_layer          = n_layer,
+        n_head           = n_head,
+        idx_len          = len(idx),
+        decoder_txt_len  = len(decoder_txt.encode("utf-8")),
+        vocab_idx_len    = len(vocab_idx),
+        vocab_txt_len    = len(vocab_txt.encode("utf-8")),
+        byte_decoder_len = len(byte_decoder),
+        )
+    with open("model.dat", "r") as f:
+        pass
+
     t2 = clock()
     print("Restore time: ", t2 - t1)
 
@@ -326,11 +364,35 @@ def convert(params,
     return mo
 
 
-def make_empty_model(blocks, n_embd, n_layer):
+def make_empty_model_with_metadata(
+        blocks           : ParamsBlocksType,
+        model_type       : int,
+        model_version    : int,
+        n_vocab          : int,
+        n_ctx            : int,
+        n_embd           : int,
+        n_layer          : int,
+        n_head           : int,
+        idx_len          : int,
+        decoder_txt_len  : int,
+        vocab_idx_len    : int,
+        vocab_txt_len    : int,
+        byte_decoder_len : int,) -> Model:
+
     mo: Model = Model(
-        blocks  = blocks,
-        n_embd  = n_embd,
-        n_layer = n_layer,
+        blocks           = blocks,
+        model_type       = model_type,
+        model_version    = model_version,
+        n_vocab          = n_vocab,
+        n_ctx            = n_ctx,
+        n_embd           = n_embd,
+        n_layer          = n_layer,
+        n_head           = n_head,
+        idx_len          = idx_len,
+        decoder_txt_len  = decoder_txt_len,
+        vocab_idx_len    = vocab_idx_len,
+        vocab_txt_len    = vocab_txt_len,
+        byte_decoder_len = byte_decoder_len,
 
         mlp_fc_w    = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32),
         mlp_fc_b    = np.empty((n_layer, 4 * n_embd)         , dtype=np.float32),
@@ -352,7 +414,20 @@ def make_empty_model(blocks, n_embd, n_layer):
     return mo
 
 
-def check_model_shapes(mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks):
+def check_model(mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks):
+    assert mo.model_type       == 0xfa51697
+    assert mo.model_version    ==       1
+    assert mo.n_vocab          ==  50_257
+    assert mo.n_ctx            ==    1024
+    assert mo.n_embd           ==     768
+    assert mo.n_layer          ==      12
+    assert mo.n_head           ==      12
+    assert mo.idx_len          ==  50_258
+    assert mo.decoder_txt_len  == 356_735
+    assert mo.vocab_idx_len    ==  50_002
+    assert mo.vocab_txt_len    == 406_304
+    assert mo.byte_decoder_len ==     256
+
     assert mo.mlp_fc_w.shape    == (nblocks,) + MlpCFcWShape
     assert mo.mlp_fc_b.shape    == (nblocks,) + MlpCFcBShape
     assert mo.mlp_proj_w.shape  == (nblocks,) + MlpCProjWShape
@@ -369,11 +444,11 @@ def check_model_shapes(mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblock
     assert mo.wpe.shape         == ParamsWpeShape
     assert mo.lnf_g.shape       == BlockLnGShape
     assert mo.lnf_b.shape       == BlockLnBShape
-    assert n_vocab              == ParamsWteShape[0] == 50_257
-    assert np.size(mo.wte, 1)   == n_embd == 768
-    assert idx.shape            == (50_258,)
-    assert vocab_idx.shape      == (50_002,)
-    assert byte_decoder.shape   == (256,)
+    assert n_vocab              == ParamsWteShape[0] == mo.n_vocab
+    assert np.size(mo.wte, 1)   == n_embd
+    assert idx.shape            == (mo.idx_len,)
+    assert vocab_idx.shape      == (mo.vocab_idx_len,)
+    assert byte_decoder.shape   == (mo.byte_decoder_len,)
 
 
 def load_decoder(filename):

From d5eb9b3e7cb3bdc4007d99340f5fd722e019f628 Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Thu, 3 Aug 2023 19:35:33 -0700
Subject: [PATCH 76/84] round-tripped metadata

---
 lp-fastgpt/create_model.py | 83 +++++++++++++++++++++++++-------------
 1 file changed, 56 insertions(+), 27 deletions(-)

diff --git a/lp-fastgpt/create_model.py b/lp-fastgpt/create_model.py
index c632e60..d4bf40b 100644
--- a/lp-fastgpt/create_model.py
+++ b/lp-fastgpt/create_model.py
@@ -183,6 +183,10 @@ def tf_load_encoder_hparams_and_params(model_size, models_dir):
                                      ParamsWteType]]
 
 
+ModelMetadataType  = np.ndarray
+ModelMetadataShape = (12,)
+
+
 @dataclass
 class Model:
     blocks      : ParamsBlocksType  # TODO: take this out of the model
@@ -302,6 +306,8 @@ def convert(params,
              len(vocab_txt.encode("utf-8")),
              len(byte_decoder)], dtype=np.int32)
 
+        assert model_metadata.shape == ModelMetadataShape
+
         model_metadata.tofile(f)
 
         mo.wte.tofile(f)
@@ -340,22 +346,41 @@ def convert(params,
 
     t1 = clock()
     m = make_empty_model_with_metadata(
-        blocks           = blocks,
-        model_type       = model_type,
-        model_version    = model_version,
-        n_vocab          = n_vocab,
-        n_ctx            = n_ctx,
-        n_embd           = n_embd,
-        n_layer          = n_layer,
-        n_head           = n_head,
-        idx_len          = len(idx),
-        decoder_txt_len  = len(decoder_txt.encode("utf-8")),
-        vocab_idx_len    = len(vocab_idx),
-        vocab_txt_len    = len(vocab_txt.encode("utf-8")),
-        byte_decoder_len = len(byte_decoder),
+        blocks           = None,  # TODO: get rid of this
+        model_type       = 0,
+        model_version    = 0,
+        n_vocab          = 0,
+        n_ctx            = 0,
+        n_embd           = 0,
+        n_layer          = 0,
+        n_head           = 0,
+        idx_len          = 0,
+        decoder_txt_len  = 0,
+        vocab_idx_len    = 0,
+        vocab_txt_len    = 0,
+        byte_decoder_len = 0,
         )
-    with open("model.dat", "r") as f:
-        pass
+
+    file_offset = 0
+    metadata = np.empty(ModelMetadataShape, dtype=np.int32)
+    metadata = np.fromfile("model.dat",
+                           dtype=np.int32,
+                           count=ModelMetadataShape[0],
+                           offset=file_offset)
+    m.model_type       = metadata[ 0]
+    m.model_version    = metadata[ 1]
+    m.n_vocab          = metadata[ 2]
+    m.n_ctx            = metadata[ 3]
+    m.n_embd           = metadata[ 4]
+    m.n_layer          = metadata[ 5]
+    m.n_head           = metadata[ 6]
+    m.idx_len          = metadata[ 7]
+    m.decoder_txt_len  = metadata[ 8]
+    m.vocab_idx_len    = metadata[ 9]
+    m.vocab_txt_len    = metadata[10]
+    m.byte_decoder_len = metadata[11]
+
+    check_model_metadata(m)
 
     t2 = clock()
     print("Restore time: ", t2 - t1)
@@ -415,18 +440,7 @@ def make_empty_model_with_metadata(
 
 
 def check_model(mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks):
-    assert mo.model_type       == 0xfa51697
-    assert mo.model_version    ==       1
-    assert mo.n_vocab          ==  50_257
-    assert mo.n_ctx            ==    1024
-    assert mo.n_embd           ==     768
-    assert mo.n_layer          ==      12
-    assert mo.n_head           ==      12
-    assert mo.idx_len          ==  50_258
-    assert mo.decoder_txt_len  == 356_735
-    assert mo.vocab_idx_len    ==  50_002
-    assert mo.vocab_txt_len    == 406_304
-    assert mo.byte_decoder_len ==     256
+    check_model_metadata(mo)
 
     assert mo.mlp_fc_w.shape    == (nblocks,) + MlpCFcWShape
     assert mo.mlp_fc_b.shape    == (nblocks,) + MlpCFcBShape
@@ -451,6 +465,21 @@ def check_model(mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks):
     assert byte_decoder.shape   == (mo.byte_decoder_len,)
 
 
+def check_model_metadata(mo):
+    assert mo.model_type       == 0xfa51697
+    assert mo.model_version    == 1
+    assert mo.n_vocab          == 50_257
+    assert mo.n_ctx            == 1024
+    assert mo.n_embd           == 768
+    assert mo.n_layer          == 12
+    assert mo.n_head           == 12
+    assert mo.idx_len          == 50_258
+    assert mo.decoder_txt_len  == 356_735
+    assert mo.vocab_idx_len    == 50_002
+    assert mo.vocab_txt_len    == 406_304
+    assert mo.byte_decoder_len == 256
+
+
 def load_decoder(filename):
     D = json.load(open(filename))
     D2 = {v: k for k, v in D.items()}

From ae88552f91b9cc6755ec25b9cf9e5689f825cfd9 Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Fri, 4 Aug 2023 07:21:32 -0700
Subject: [PATCH 77/84] round-tripped model data -- indices and texts left to
 do

---
 lp-fastgpt/create_model.py | 235 +++++++++++++++++++++++++++----------
 1 file changed, 171 insertions(+), 64 deletions(-)

diff --git a/lp-fastgpt/create_model.py b/lp-fastgpt/create_model.py
index d4bf40b..75c4ae7 100644
--- a/lp-fastgpt/create_model.py
+++ b/lp-fastgpt/create_model.py
@@ -36,6 +36,7 @@
 
 """
 
+
 from time import monotonic as clock
 import os
 import json
@@ -128,40 +129,70 @@ def tf_load_encoder_hparams_and_params(model_size, models_dir):
 
 from dataclasses import dataclass
 
-HParamsType = dict[str, int]
+
+# === Magic (Unexplained) Numbers =========================
+
+DecoderIdxType     = np.ndarray
+DecoderIdxShape    = (50_258,)
+
+DecoderTxtType     = str
+DecoderTxtAsciiLen = 320_827
+DecoderTxtUtf8Len  = 356_735  # It's checked somewhere below.
+
+VocabIdxType       = np.ndarray
+VocabIdxShape      = (50_002,)
+
+VocabTxtType       = str
+VocabTxtAsciiLen   = 370_558
+VocabTxtUtf8Len    = 406_304  # It's checked somewhere below.
+NVocab             = 50_257
+
+DecoderShape       = (256,)
+
+NBlocks            = 12
+
+ModelType          = 0xfa51697
+ModelVersion       = 1
+
+NCtx               = 1024
+NEmbed             = 768
+
+Magic2304          = 2304
+Magic3072          = 3072
+
 
 # === Block-Attention components ======================
 CAttnBType         = np.ndarray
-CAttnBShape        = (     2304,)
+CAttnBShape        = (        Magic2304,)
 CAttnWType         = np.ndarray
-CAttnWShape        = (768, 2304,)
+CAttnWShape        = (NEmbed, Magic2304,)
 CAttnType          = dict[str, Union[CAttnBType, CAttnWType]]
 
 CProjBType         = np.ndarray
-CProjBShape        = (      768,)
+CProjBShape        = (        NEmbed,)
 CProjWType         = np.ndarray
-CProjWShape        = (768,  768,)
+CProjWShape        = (NEmbed, NEmbed,)
 CProjType          = dict[str, Union[CProjBType, CProjWType]]
 # --- top level
 BlockAttnType      = dict[str, Union[CAttnType, CProjType]]
-# === BlockIn components ==============================
+# === BlockLn components ==============================
 BlockLnBType       = np.ndarray
-BlockLnBShape      = (768,)
+BlockLnBShape      = (NEmbed,)
 BlockLnGType       = np.ndarray
-BlockLnGShape      = (768,)
+BlockLnGShape      = (NEmbed,)
 # --- top level
 BlockLnType        = dict[str, Union[BlockLnBType, BlockLnGType]]
 # === BlockMlp components =============================
 MlpCFcBType        = np.ndarray
-MlpCFcBShape       = (     3072,)
+MlpCFcBShape       = (        Magic3072,)
 MlpCFcWType        = np.ndarray
-MlpCFcWShape       = (768, 3072,)
+MlpCFcWShape       = (NEmbed, Magic3072,)
 MlpCFcType         = dict[str, Union[MlpCFcBType, MlpCFcWType]]
 
 MlpCProjBType      = np.ndarray
-MlpCProjBShape     = (      768,)
+MlpCProjBShape     = (          NEmbed,)
 MlpCProjWType      = np.ndarray
-MlpCProjWShape     = (3072, 768,)
+MlpCProjWShape     = (Magic3072, NEmbed,)
 MlpCProjType       = dict[str, Union[MlpCProjBType, MlpCProjWType]]
 # --- top level
 BlockMlpType       = dict[str, Union[MlpCFcType, MlpCProjType]]
@@ -169,27 +200,27 @@ def tf_load_encoder_hparams_and_params(model_size, models_dir):
 ParamsBlockType    = dict[str, Union[BlockAttnType,
                                      BlockLnType,  # two of these
                                      BlockMlpType]]
-#                 len=12
+
 ParamsBlocksType   = list[ParamsBlockType]
 ParamsLnFType      = dict[str, np.ndarray]
 ParamsLnFValShape  = BlockLnBShape
 ParamsWpeType      = np.ndarray
-ParamsWpeShape     = ( 1024, 768,)
+ParamsWpeShape     = (  NCtx, NEmbed,)
 ParamsWteType      = np.ndarray
-ParamsWteShape     = (50257, 768,)
+ParamsWteShape     = (NVocab, NEmbed,)
 ParamsType         = dict[str, Union[ParamsBlocksType,
                                      ParamsLnFType,
                                      ParamsWpeType,
                                      ParamsWteType]]
 
-
 ModelMetadataType  = np.ndarray
 ModelMetadataShape = (12,)
 
+HParamsType = dict[str, int]
+
 
 @dataclass
 class Model:
-    blocks      : ParamsBlocksType  # TODO: take this out of the model
     # integer metadata
     model_type       : int
     model_version    : int
@@ -198,7 +229,7 @@ class Model:
     n_embd           : int
     n_layer          : int
     n_head           : int
-    idx_len          : int
+    decoder_idx_len  : int
     decoder_txt_len  : int
     vocab_idx_len    : int
     vocab_txt_len    : int
@@ -225,7 +256,7 @@ class Model:
 def convert(params,
             n_head,
             n_ctx,
-            idx,
+            decoder_idx,
             decoder_txt,
             vocab_idx,
             vocab_txt,
@@ -236,18 +267,17 @@ def convert(params,
     # must predefine just to get shapes ...
     blocks : ParamsBlocksType = params["blocks"]
     nblocks = len(blocks)
-    assert nblocks == 12
+    assert nblocks == NBlocks
 
     n_embd  = blocks[0]["ln_1"]["b"].size
     n_layer = nblocks
-    assert n_layer == 12
+    assert n_layer == NBlocks
 
-    n_vocab = ParamsWteShape[0]  # np.size(mo.wte, 0)
-    model_type = 0xfa51697  # fastGPT
-    model_version = 1
+    n_vocab       = ParamsWteShape[0]  # np.size(mo.wte, 0)
+    model_type    = ModelType
+    model_version = ModelVersion
 
     mo : Model = make_empty_model_with_metadata(
-        blocks           = blocks,
         model_type       = model_type,
         model_version    = model_version,
         n_vocab          = n_vocab,
@@ -255,9 +285,10 @@ def convert(params,
         n_embd           = n_embd,
         n_layer          = n_layer,
         n_head           = n_head,
-        idx_len          = len(idx),
+        decoder_idx_len  = DecoderIdxShape[0],
+        # It's useful to check magic numbers against computations:
         decoder_txt_len  = len(decoder_txt.encode("utf-8")),
-        vocab_idx_len    = len(vocab_idx),
+        vocab_idx_len    = VocabIdxShape[0],
         vocab_txt_len    = len(vocab_txt.encode("utf-8")),
         byte_decoder_len = len(byte_decoder),
     )
@@ -279,6 +310,7 @@ def convert(params,
 
     mo.wte   = params["wte"]
     mo.wpe   = params["wpe"]
+
     mo.lnf_g = params["ln_f"]["g"]
     mo.lnf_b = params["ln_f"]["b"]
 
@@ -286,7 +318,7 @@ def convert(params,
     print("Transform time: ", t2 - t1)
 
     check_model(
-        mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks)
+        mo, n_vocab, decoder_idx, vocab_idx, byte_decoder, n_embd, nblocks)
 
     # Save the model
     t1 = clock()
@@ -300,7 +332,7 @@ def convert(params,
              n_embd,
              n_layer,
              n_head,
-             len(idx),
+             DecoderIdxShape[0],
              len(decoder_txt.encode("utf-8")),
              len(vocab_idx),
              len(vocab_txt.encode("utf-8")),
@@ -329,10 +361,12 @@ def convert(params,
         mo.ln1_g.tofile(f)
         mo.ln2_b.tofile(f)
         mo.ln2_g.tofile(f)
+
         mo.lnf_b.tofile(f)
         mo.lnf_g.tofile(f)
 
-        idx.tofile(f)
+        decoder_idx.tofile(f)
+
         f.write(decoder_txt)
 
         vocab_idx.tofile(f)
@@ -346,7 +380,6 @@ def convert(params,
 
     t1 = clock()
     m = make_empty_model_with_metadata(
-        blocks           = None,  # TODO: get rid of this
         model_type       = 0,
         model_version    = 0,
         n_vocab          = 0,
@@ -354,19 +387,18 @@ def convert(params,
         n_embd           = 0,
         n_layer          = 0,
         n_head           = 0,
-        idx_len          = 0,
+        decoder_idx_len  = 0,
         decoder_txt_len  = 0,
         vocab_idx_len    = 0,
         vocab_txt_len    = 0,
         byte_decoder_len = 0,
         )
 
-    file_offset = 0
-    metadata = np.empty(ModelMetadataShape, dtype=np.int32)
+    floff = 0
     metadata = np.fromfile("model.dat",
                            dtype=np.int32,
                            count=ModelMetadataShape[0],
-                           offset=file_offset)
+                           offset=floff)
     m.model_type       = metadata[ 0]
     m.model_version    = metadata[ 1]
     m.n_vocab          = metadata[ 2]
@@ -374,7 +406,7 @@ def convert(params,
     m.n_embd           = metadata[ 4]
     m.n_layer          = metadata[ 5]
     m.n_head           = metadata[ 6]
-    m.idx_len          = metadata[ 7]
+    m.decoder_idx_len  = metadata[ 7]
     m.decoder_txt_len  = metadata[ 8]
     m.vocab_idx_len    = metadata[ 9]
     m.vocab_txt_len    = metadata[10]
@@ -382,15 +414,90 @@ def convert(params,
 
     check_model_metadata(m)
 
+    floff += ModelMetadataShape[0] * BYTES_PER_INT32
+
+    floff, m.wte = restore_floats(ParamsWteShape, floff)
+    assert np.all(m.wte == mo.wte)
+
+    floff, m.wpe = restore_floats(ParamsWpeShape, floff)
+    assert np.all(m.wpe == mo.wpe)
+
+    floff, m.mlp_fc_w = restore_floats((nblocks,) + MlpCFcWShape, floff)
+    assert np.all(m.mlp_fc_w == mo.mlp_fc_w)
+
+    floff, m.mlp_fc_b = restore_floats((nblocks,) + MlpCFcBShape, floff)
+    assert np.all(m.mlp_fc_b == mo.mlp_fc_b)
+
+    floff, m.mlp_proj_w = restore_floats((nblocks,) + MlpCProjWShape, floff)
+    assert np.all(m.mlp_proj_w == mo.mlp_proj_w)
+
+    floff, m.mlp_proj_b = restore_floats((nblocks,) + MlpCProjBShape, floff)
+    assert np.all(m.mlp_proj_b == mo.mlp_proj_b)
+
+    floff, m.attn_w = restore_floats((nblocks,) + CAttnWShape, floff)
+    assert np.all(m.attn_w == mo.attn_w)
+
+    floff, m.attn_b = restore_floats((nblocks,) + CAttnBShape, floff)
+    assert np.all(m.attn_b == mo.attn_b)
+
+    floff, m.attn_proj_w = restore_floats((nblocks,) + CProjWShape, floff)
+    assert np.all(m.attn_proj_w == mo.attn_proj_w)
+
+    floff, m.attn_proj_b = restore_floats((nblocks,) + CProjBShape, floff)
+    assert np.all(m.attn_proj_b == mo.attn_proj_b)
+
+    floff, m.ln1_b = restore_floats((nblocks,) + BlockLnBShape, floff)
+    assert np.all(m.ln1_b == mo.ln1_b)
+
+    floff, m.ln1_g = restore_floats((nblocks,) + BlockLnBShape, floff)
+    assert np.all(m.ln1_g == mo.ln1_g)
+
+    floff, m.ln2_b = restore_floats((nblocks,) + BlockLnBShape, floff)
+    assert np.all(m.ln2_b == mo.ln2_b)
+
+    floff, m.ln2_g = restore_floats((nblocks,) + BlockLnBShape, floff)
+    assert np.all(m.ln2_g == mo.ln2_g)
+
+    floff, m.lnf_b = restore_floats(ParamsLnFValShape, floff)
+    assert np.all(m.lnf_b == mo.lnf_b)
+
+    floff, m.lnf_g = restore_floats(ParamsLnFValShape, floff)
+    assert np.all(m.lnf_g == mo.lnf_g)
+
+    # floff, decoder_idxi = restore_floats(DecoderIdxShape, floff)
+    # assert np.all(decoder_idxi == decoder_idx)
+
     t2 = clock()
     print("Restore time: ", t2 - t1)
 
-
     return mo
 
 
+def prod_tuple(t : tuple[int]) -> int:
+    result : int = 1
+    for e in t:
+        result *= e
+    return result
+
+
+BYTES_PER_INT32   = 4
+BYTES_PER_FLOAT32 = 4
+
+
+def restore_floats(shape : tuple, offset : int) -> tuple[int, np.ndarray]:
+    """agnostic to length of shape; TODO impossible to statically type"""
+    result : np.ndarray
+    count = prod_tuple(shape)
+    result = np.fromfile("model.dat",
+                         dtype=np.float32,
+                         count=count,
+                         offset=offset)
+    result = np.reshape(result, shape)
+    new_offset : int = offset + (count * BYTES_PER_FLOAT32)
+    return new_offset, result
+
+
 def make_empty_model_with_metadata(
-        blocks           : ParamsBlocksType,
         model_type       : int,
         model_version    : int,
         n_vocab          : int,
@@ -398,14 +505,13 @@ def make_empty_model_with_metadata(
         n_embd           : int,
         n_layer          : int,
         n_head           : int,
-        idx_len          : int,
+        decoder_idx_len  : int,
         decoder_txt_len  : int,
         vocab_idx_len    : int,
         vocab_txt_len    : int,
         byte_decoder_len : int,) -> Model:
 
     mo: Model = Model(
-        blocks           = blocks,
         model_type       = model_type,
         model_version    = model_version,
         n_vocab          = n_vocab,
@@ -413,7 +519,7 @@ def make_empty_model_with_metadata(
         n_embd           = n_embd,
         n_layer          = n_layer,
         n_head           = n_head,
-        idx_len          = idx_len,
+        decoder_idx_len  = decoder_idx_len,
         decoder_txt_len  = decoder_txt_len,
         vocab_idx_len    = vocab_idx_len,
         vocab_txt_len    = vocab_txt_len,
@@ -460,24 +566,24 @@ def check_model(mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks):
     assert mo.lnf_b.shape       == BlockLnBShape
     assert n_vocab              == ParamsWteShape[0] == mo.n_vocab
     assert np.size(mo.wte, 1)   == n_embd
-    assert idx.shape            == (mo.idx_len,)
+    assert idx.shape            == (mo.decoder_idx_len,)
     assert vocab_idx.shape      == (mo.vocab_idx_len,)
     assert byte_decoder.shape   == (mo.byte_decoder_len,)
 
 
 def check_model_metadata(mo):
-    assert mo.model_type       == 0xfa51697
-    assert mo.model_version    == 1
-    assert mo.n_vocab          == 50_257
-    assert mo.n_ctx            == 1024
-    assert mo.n_embd           == 768
-    assert mo.n_layer          == 12
-    assert mo.n_head           == 12
-    assert mo.idx_len          == 50_258
-    assert mo.decoder_txt_len  == 356_735
-    assert mo.vocab_idx_len    == 50_002
-    assert mo.vocab_txt_len    == 406_304
-    assert mo.byte_decoder_len == 256
+    assert mo.model_type       == ModelType
+    assert mo.model_version    == ModelVersion
+    assert mo.n_vocab          == NVocab
+    assert mo.n_ctx            == NCtx
+    assert mo.n_embd           == NEmbed
+    assert mo.n_layer          == NBlocks
+    assert mo.n_head           == NBlocks
+    assert mo.decoder_idx_len  == DecoderIdxShape[0]
+    assert mo.decoder_txt_len  == DecoderTxtUtf8Len
+    assert mo.vocab_idx_len    == VocabIdxShape[0]
+    assert mo.vocab_txt_len    == VocabTxtUtf8Len
+    assert mo.byte_decoder_len == DecoderShape[0]
 
 
 def load_decoder(filename):
@@ -500,7 +606,7 @@ def load_vocab(filename):
     return D
 
 
-def decoder_idx(decoder):
+def load_decoder_idx(decoder):
     i = 0
     idx = np.empty(len(decoder) + 1, dtype=np.int32)
     idx[0] = i
@@ -529,7 +635,7 @@ def bytes_to_unicode():
     for y in byte_decoder:
         x = ord(y)
         bd[x] = byte_decoder[y]
-    bd2 = np.zeros(256, dtype=np.int32)
+    bd2 = np.zeros(DecoderShape, dtype=np.int32)
     for i in range(np.size(bd)):
         bd2[bd[i]] = i
     return bd2
@@ -552,7 +658,7 @@ def main(model_size: str = "124M",
     decoder : list[str] = \
         load_decoder(os.path.join(models_dir, model_size, "encoder.json"))
 
-    assert len(decoder) == 50_257
+    assert len(decoder) == NVocab  # TODO: ??? !!! ???
 
     vocab : list[str] = \
         load_vocab(os.path.join(models_dir, model_size, "vocab.bpe"))
@@ -566,19 +672,20 @@ def main(model_size: str = "124M",
     print("Converting model, saving to `model.dat`")
     t1 = clock()
 
+    idx : np.ndarray = load_decoder_idx(decoder)
+    assert idx.shape == DecoderIdxShape
+
     decoder_txt : str = "".join(decoder)
+    assert len(decoder_txt) == DecoderTxtAsciiLen
 
-    idx : np.ndarray = decoder_idx(decoder)
+    vocab_idx    = load_decoder_idx(vocab)
+    assert vocab_idx.shape == VocabIdxShape
 
     vocab_txt    = "".join(vocab)
-
-    vocab_idx    = decoder_idx(vocab)
-
-    assert vocab_idx.shape == (50_002,)
+    assert len(vocab_txt) == VocabTxtAsciiLen
 
     byte_decoder = bytes_to_unicode()
-
-    assert byte_decoder.shape == (256,)
+    assert byte_decoder.shape == DecoderShape
 
     m : Model = \
         convert(params,

From d7405f05434ff113b932e3296f85f853ed10e60f Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Fri, 4 Aug 2023 16:48:41 -0700
Subject: [PATCH 78/84] everything round-tripped

---
 lp-fastgpt/create_model.py | 40 +++++++++++++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

diff --git a/lp-fastgpt/create_model.py b/lp-fastgpt/create_model.py
index 75c4ae7..9340437 100644
--- a/lp-fastgpt/create_model.py
+++ b/lp-fastgpt/create_model.py
@@ -148,6 +148,7 @@ def tf_load_encoder_hparams_and_params(model_size, models_dir):
 NVocab             = 50_257
 
 DecoderShape       = (256,)
+DecoderLen         = DecoderShape[0]
 
 NBlocks            = 12
 
@@ -464,8 +465,28 @@ def convert(params,
     floff, m.lnf_g = restore_floats(ParamsLnFValShape, floff)
     assert np.all(m.lnf_g == mo.lnf_g)
 
-    # floff, decoder_idxi = restore_floats(DecoderIdxShape, floff)
-    # assert np.all(decoder_idxi == decoder_idx)
+    floff, decoder_idxi = restore_ints(DecoderIdxShape, floff)
+    assert np.all(decoder_idxi == decoder_idx)
+
+    with open("model.dat", "rb") as f:
+        f.seek(floff)
+        decoder_txti_ub : bytes = f.read(DecoderTxtUtf8Len)
+        decoder_txti_u = decoder_txti_ub.decode("utf-8")
+        assert len(decoder_txti_u) == DecoderTxtAsciiLen
+        floff += DecoderTxtUtf8Len
+
+    floff, vocab_idxi = restore_ints(VocabIdxShape, floff)
+    assert np.all(vocab_idxi == vocab_idx)
+
+    with open("model.dat", "rb") as f:
+        f.seek(floff)
+        vocab_txt_ub : bytes = f.read(VocabTxtUtf8Len)
+        vocab_txt_u = vocab_txt_ub.decode("utf-8")
+        assert len(vocab_txt_u) == VocabTxtAsciiLen
+        floff += VocabTxtUtf8Len
+
+    floff, byte_decoderi = restore_ints(DecoderShape, floff)
+    assert np.all(vocab_idxi == vocab_idx)
 
     t2 = clock()
     print("Restore time: ", t2 - t1)
@@ -497,6 +518,19 @@ def restore_floats(shape : tuple, offset : int) -> tuple[int, np.ndarray]:
     return new_offset, result
 
 
+def restore_ints(shape : tuple, offset : int) -> tuple[int, np.ndarray]:
+    """agnostic to length of shape; TODO impossible to statically type"""
+    result : np.ndarray
+    count = prod_tuple(shape)
+    result = np.fromfile("model.dat",
+                         dtype=np.int32,
+                         count=count,
+                         offset=offset)
+    result = np.reshape(result, shape)
+    new_offset : int = offset + (count * BYTES_PER_INT32)
+    return new_offset, result
+
+
 def make_empty_model_with_metadata(
         model_type       : int,
         model_version    : int,
@@ -617,7 +651,7 @@ def load_decoder_idx(decoder):
     return idx
 
 
-def bytes_to_unicode():
+def bytes_to_unicode() -> np.ndarray:
     bs = list(range(ord("!"), ord("~") + 1)) + \
         list(range(ord("¡"), ord("¬") + 1)) + \
         list(range(ord("®"), ord("ÿ") + 1))

From 495a532bc226aee256acd20703369e66362f3526 Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Fri, 4 Aug 2023 19:37:12 -0700
Subject: [PATCH 79/84] nearly done with round-tripping

---
 lp-fastgpt/create_model.py | 104 ++++++++++++++++++++++++-------------
 1 file changed, 68 insertions(+), 36 deletions(-)

diff --git a/lp-fastgpt/create_model.py b/lp-fastgpt/create_model.py
index 9340437..d86289a 100644
--- a/lp-fastgpt/create_model.py
+++ b/lp-fastgpt/create_model.py
@@ -236,22 +236,28 @@ class Model:
     vocab_txt_len    : int
     byte_decoder_len : int
     # check shapes in asserts for now; type system too weak.
-    mlp_fc_w    : np.ndarray
-    mlp_fc_b    : np.ndarray
-    mlp_proj_w  : np.ndarray
-    mlp_proj_b  : np.ndarray
-    attn_w      : np.ndarray
-    attn_b      : np.ndarray
-    attn_proj_w : np.ndarray
-    attn_proj_b : np.ndarray
-    ln1_g       : np.ndarray
-    ln1_b       : np.ndarray
-    ln2_g       : np.ndarray
-    ln2_b       : np.ndarray
-    wte         : np.ndarray
-    wpe         : np.ndarray
-    lnf_g       : np.ndarray
-    lnf_b       : np.ndarray
+    mlp_fc_w         : np.ndarray
+    mlp_fc_b         : np.ndarray
+    mlp_proj_w       : np.ndarray
+    mlp_proj_b       : np.ndarray
+    attn_w           : np.ndarray
+    attn_b           : np.ndarray
+    attn_proj_w      : np.ndarray
+    attn_proj_b      : np.ndarray
+    ln1_g            : np.ndarray
+    ln1_b            : np.ndarray
+    ln2_g            : np.ndarray
+    ln2_b            : np.ndarray
+    wte              : np.ndarray
+    wpe              : np.ndarray
+    lnf_g            : np.ndarray
+    lnf_b            : np.ndarray
+    # auxiliary matrices and texts
+    decoder_idx      : np.ndarray
+    decoder_txt      : str
+    vocab_idx        : np.ndarray
+    vocab_txt        : str
+    byte_decoder     : np.ndarray
 
 
 def convert(params,
@@ -315,6 +321,12 @@ def convert(params,
     mo.lnf_g = params["ln_f"]["g"]
     mo.lnf_b = params["ln_f"]["b"]
 
+    mo.decoder_idx  = decoder_idx
+    mo.decoder_txt  = decoder_txt
+    mo.vocab_idx    = vocab_idx
+    mo.vocab_txt    = vocab_txt
+    mo.byte_decoder = byte_decoder
+
     t2 = clock()
     print("Transform time: ", t2 - t1)
 
@@ -467,26 +479,33 @@ def convert(params,
 
     floff, decoder_idxi = restore_ints(DecoderIdxShape, floff)
     assert np.all(decoder_idxi == decoder_idx)
+    m.decoder_idx = decoder_idxi
 
     with open("model.dat", "rb") as f:
         f.seek(floff)
         decoder_txti_ub : bytes = f.read(DecoderTxtUtf8Len)
         decoder_txti_u = decoder_txti_ub.decode("utf-8")
         assert len(decoder_txti_u) == DecoderTxtAsciiLen
+        assert decoder_txti_u == decoder_txt
+        m.decoder_idx = decoder_txti_u
         floff += DecoderTxtUtf8Len
 
     floff, vocab_idxi = restore_ints(VocabIdxShape, floff)
     assert np.all(vocab_idxi == vocab_idx)
+    m.vocab_idx = vocab_idxi
 
     with open("model.dat", "rb") as f:
         f.seek(floff)
         vocab_txt_ub : bytes = f.read(VocabTxtUtf8Len)
         vocab_txt_u = vocab_txt_ub.decode("utf-8")
         assert len(vocab_txt_u) == VocabTxtAsciiLen
+        assert vocab_txt_u == vocab_txt
+        m.vocab_txt = vocab_txt_u
         floff += VocabTxtUtf8Len
 
     floff, byte_decoderi = restore_ints(DecoderShape, floff)
-    assert np.all(vocab_idxi == vocab_idx)
+    assert np.all(byte_decoderi == byte_decoder)
+    m.byte_decoder = byte_decoderi
 
     t2 = clock()
     print("Restore time: ", t2 - t1)
@@ -559,27 +578,39 @@ def make_empty_model_with_metadata(
         vocab_txt_len    = vocab_txt_len,
         byte_decoder_len = byte_decoder_len,
 
-        mlp_fc_w    = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32),
-        mlp_fc_b    = np.empty((n_layer, 4 * n_embd)         , dtype=np.float32),
-        mlp_proj_w  = np.empty((n_layer, 4 * n_embd, n_embd) , dtype=np.float32),
-        mlp_proj_b  = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        attn_w      = np.empty((n_layer, n_embd, 3 * n_embd) , dtype=np.float32),
-        attn_b      = np.empty((n_layer, 3 * n_embd)         , dtype=np.float32),
-        attn_proj_w = np.empty((n_layer, n_embd, n_embd)     , dtype=np.float32),
-        attn_proj_b = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        ln1_g       = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        ln1_b       = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        ln2_g       = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        ln2_b       = np.empty((n_layer, n_embd)             , dtype=np.float32),
-        wte         = np.empty(0                             , dtype=np.float32),
-        wpe         = np.empty(0                             , dtype=np.float32),
-        lnf_g       = np.empty(0                             , dtype=np.float32),
-        lnf_b       = np.empty(0                             , dtype=np.float32),
+        mlp_fc_w     = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32),
+        mlp_fc_b     = np.empty((n_layer, 4 * n_embd)         , dtype=np.float32),
+        mlp_proj_w   = np.empty((n_layer, 4 * n_embd, n_embd) , dtype=np.float32),
+        mlp_proj_b   = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        attn_w       = np.empty((n_layer, n_embd, 3 * n_embd) , dtype=np.float32),
+        attn_b       = np.empty((n_layer, 3 * n_embd)         , dtype=np.float32),
+        attn_proj_w  = np.empty((n_layer, n_embd, n_embd)     , dtype=np.float32),
+        attn_proj_b  = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln1_g        = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln1_b        = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln2_g        = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln2_b        = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        wte          = np.empty(0                             , dtype=np.float32),
+        wpe          = np.empty(0                             , dtype=np.float32),
+        lnf_g        = np.empty(0                             , dtype=np.float32),
+        lnf_b        = np.empty(0                             , dtype=np.float32),
+
+        decoder_idx  = np.empty(0                             , dtype=np.int32),
+        decoder_txt  = '',
+        vocab_idx    = np.empty(0                             , dtype=np.int32),
+        vocab_txt    = '',
+        byte_decoder = np.empty(0                             , dtype=np.int32),
     )
     return mo
 
 
-def check_model(mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks):
+def check_model(mo,
+                n_vocab,
+                decoder_idx,
+                vocab_idx,
+                byte_decoder,
+                n_embd,
+                nblocks):
     check_model_metadata(mo)
 
     assert mo.mlp_fc_w.shape    == (nblocks,) + MlpCFcWShape
@@ -598,13 +629,13 @@ def check_model(mo, n_vocab, idx, vocab_idx, byte_decoder, n_embd, nblocks):
     assert mo.wpe.shape         == ParamsWpeShape
     assert mo.lnf_g.shape       == BlockLnGShape
     assert mo.lnf_b.shape       == BlockLnBShape
+
     assert n_vocab              == ParamsWteShape[0] == mo.n_vocab
     assert np.size(mo.wte, 1)   == n_embd
-    assert idx.shape            == (mo.decoder_idx_len,)
+    assert decoder_idx.shape    == (mo.decoder_idx_len,)
     assert vocab_idx.shape      == (mo.vocab_idx_len,)
     assert byte_decoder.shape   == (mo.byte_decoder_len,)
 
-
 def check_model_metadata(mo):
     assert mo.model_type       == ModelType
     assert mo.model_version    == ModelVersion
@@ -613,6 +644,7 @@ def check_model_metadata(mo):
     assert mo.n_embd           == NEmbed
     assert mo.n_layer          == NBlocks
     assert mo.n_head           == NBlocks
+
     assert mo.decoder_idx_len  == DecoderIdxShape[0]
     assert mo.decoder_txt_len  == DecoderTxtUtf8Len
     assert mo.vocab_idx_len    == VocabIdxShape[0]

From 7f519060b29c3ebb632692b842522aa01897b92e Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Sat, 5 Aug 2023 06:19:08 -0700
Subject: [PATCH 80/84] removed TF dependencies from 'restore_model'

---
 lp-fastgpt/create_model.py  | 252 +++++++++++-----------
 lp-fastgpt/main.py          |  19 +-
 lp-fastgpt/restore_model.py | 402 ++++++++++++++++++++++++++++++++++++
 3 files changed, 521 insertions(+), 152 deletions(-)
 create mode 100644 lp-fastgpt/restore_model.py

diff --git a/lp-fastgpt/create_model.py b/lp-fastgpt/create_model.py
index d86289a..0d3cb83 100644
--- a/lp-fastgpt/create_model.py
+++ b/lp-fastgpt/create_model.py
@@ -330,26 +330,64 @@ def convert(params,
     t2 = clock()
     print("Transform time: ", t2 - t1)
 
-    check_model(
-        mo, n_vocab, decoder_idx, vocab_idx, byte_decoder, n_embd, nblocks)
+    check_model(mo)
 
     # Save the model
     t1 = clock()
-    with open("model.dat", "w") as f:
+    save_model(mo)
+    t2 = clock()
+    print("Save time: ", t2 - t1)
+
+    # Round-trip the model: check the stored version.
+
+    t1 = clock()
+    m = restore_model()
+
+    check_model(m)
+    assert np.all(m.wte          == mo.wte)
+    assert np.all(m.wpe          == mo.wpe)
+    assert np.all(m.mlp_fc_w     == mo.mlp_fc_w)
+    assert np.all(m.mlp_fc_b     == mo.mlp_fc_b)
+    assert np.all(m.mlp_proj_w   == mo.mlp_proj_w)
+    assert np.all(m.mlp_proj_b   == mo.mlp_proj_b)
+    assert np.all(m.attn_w       == mo.attn_w)
+    assert np.all(m.attn_b       == mo.attn_b)
+    assert np.all(m.attn_proj_w  == mo.attn_proj_w)
+    assert np.all(m.attn_proj_b  == mo.attn_proj_b)
+    assert np.all(m.ln1_b        == mo.ln1_b)
+    assert np.all(m.ln1_g        == mo.ln1_g)
+    assert np.all(m.ln2_b        == mo.ln2_b)
+    assert np.all(m.ln2_g        == mo.ln2_g)
+    assert np.all(m.lnf_b        == mo.lnf_b)
+    assert np.all(m.lnf_g        == mo.lnf_g)
+    assert np.all(m.decoder_idx  == mo.decoder_idx)
+    assert m.decoder_txt         == mo.decoder_txt
+    assert np.all(m.vocab_idx    == mo.vocab_idx)
+    assert m.vocab_txt           == mo.vocab_txt
+    assert np.all(m.byte_decoder == mo.byte_decoder)
+    t2 = clock()
+    print("Restore time: ", t2 - t1)
+
+    return m
+
 
+def save_model(mo : Model):
+    with open("model.dat", "w") as f:
         model_metadata = np.array(
-            [model_type,
-             model_version,
-             n_vocab,
-             n_ctx,
-             n_embd,
-             n_layer,
-             n_head,
-             DecoderIdxShape[0],
-             len(decoder_txt.encode("utf-8")),
-             len(vocab_idx),
-             len(vocab_txt.encode("utf-8")),
-             len(byte_decoder)], dtype=np.int32)
+            [
+                mo.model_type,
+                mo.model_version,
+                mo.n_vocab,
+                mo.n_ctx,
+                mo.n_embd,
+                mo.n_layer,
+                mo.n_head,
+                len(mo.decoder_idx),
+                len(mo.decoder_txt.encode("utf-8")),
+                len(mo.vocab_idx),
+                len(mo.vocab_txt.encode("utf-8")),
+                len(mo.byte_decoder)],
+            dtype=np.int32)
 
         assert model_metadata.shape == ModelMetadataShape
 
@@ -378,21 +416,15 @@ def convert(params,
         mo.lnf_b.tofile(f)
         mo.lnf_g.tofile(f)
 
-        decoder_idx.tofile(f)
+        mo.decoder_idx.tofile(f)
+        f.write(mo.decoder_txt)
+        mo.vocab_idx.tofile(f)
+        f.write(mo.vocab_txt)
+        mo.byte_decoder.tofile(f)
 
-        f.write(decoder_txt)
 
-        vocab_idx.tofile(f)
-
-        f.write(vocab_txt)
-
-        byte_decoder.tofile(f)
-
-    t2 = clock()
-    print("Save time: ", t2 - t1)
-
-    t1 = clock()
-    m = make_empty_model_with_metadata(
+def restore_model() -> Model:
+    m : Model = make_empty_model_with_metadata(
         model_type       = 0,
         model_version    = 0,
         n_vocab          = 0,
@@ -407,11 +439,13 @@ def convert(params,
         byte_decoder_len = 0,
         )
 
-    floff = 0
-    metadata = np.fromfile("model.dat",
-                           dtype=np.int32,
-                           count=ModelMetadataShape[0],
-                           offset=floff)
+    floff : int = 0
+    metadata : np.ndarray = (
+        np.fromfile("model.dat",
+                    dtype=np.int32,
+                    count=ModelMetadataShape[0],
+                    offset=floff))
+
     m.model_type       = metadata[ 0]
     m.model_version    = metadata[ 1]
     m.n_vocab          = metadata[ 2]
@@ -428,89 +462,43 @@ def convert(params,
     check_model_metadata(m)
 
     floff += ModelMetadataShape[0] * BYTES_PER_INT32
-
-    floff, m.wte = restore_floats(ParamsWteShape, floff)
-    assert np.all(m.wte == mo.wte)
-
-    floff, m.wpe = restore_floats(ParamsWpeShape, floff)
-    assert np.all(m.wpe == mo.wpe)
-
-    floff, m.mlp_fc_w = restore_floats((nblocks,) + MlpCFcWShape, floff)
-    assert np.all(m.mlp_fc_w == mo.mlp_fc_w)
-
-    floff, m.mlp_fc_b = restore_floats((nblocks,) + MlpCFcBShape, floff)
-    assert np.all(m.mlp_fc_b == mo.mlp_fc_b)
-
-    floff, m.mlp_proj_w = restore_floats((nblocks,) + MlpCProjWShape, floff)
-    assert np.all(m.mlp_proj_w == mo.mlp_proj_w)
-
-    floff, m.mlp_proj_b = restore_floats((nblocks,) + MlpCProjBShape, floff)
-    assert np.all(m.mlp_proj_b == mo.mlp_proj_b)
-
-    floff, m.attn_w = restore_floats((nblocks,) + CAttnWShape, floff)
-    assert np.all(m.attn_w == mo.attn_w)
-
-    floff, m.attn_b = restore_floats((nblocks,) + CAttnBShape, floff)
-    assert np.all(m.attn_b == mo.attn_b)
-
-    floff, m.attn_proj_w = restore_floats((nblocks,) + CProjWShape, floff)
-    assert np.all(m.attn_proj_w == mo.attn_proj_w)
-
-    floff, m.attn_proj_b = restore_floats((nblocks,) + CProjBShape, floff)
-    assert np.all(m.attn_proj_b == mo.attn_proj_b)
-
-    floff, m.ln1_b = restore_floats((nblocks,) + BlockLnBShape, floff)
-    assert np.all(m.ln1_b == mo.ln1_b)
-
-    floff, m.ln1_g = restore_floats((nblocks,) + BlockLnBShape, floff)
-    assert np.all(m.ln1_g == mo.ln1_g)
-
-    floff, m.ln2_b = restore_floats((nblocks,) + BlockLnBShape, floff)
-    assert np.all(m.ln2_b == mo.ln2_b)
-
-    floff, m.ln2_g = restore_floats((nblocks,) + BlockLnBShape, floff)
-    assert np.all(m.ln2_g == mo.ln2_g)
-
-    floff, m.lnf_b = restore_floats(ParamsLnFValShape, floff)
-    assert np.all(m.lnf_b == mo.lnf_b)
-
-    floff, m.lnf_g = restore_floats(ParamsLnFValShape, floff)
-    assert np.all(m.lnf_g == mo.lnf_g)
-
-    floff, decoder_idxi = restore_ints(DecoderIdxShape, floff)
-    assert np.all(decoder_idxi == decoder_idx)
-    m.decoder_idx = decoder_idxi
+    floff, m.wte         = restore_floats(ParamsWteShape              , floff)
+    floff, m.wpe         = restore_floats(ParamsWpeShape              , floff)
+    floff, m.mlp_fc_w    = restore_floats((NBlocks,) + MlpCFcWShape   , floff)
+    floff, m.mlp_fc_b    = restore_floats((NBlocks,) + MlpCFcBShape   , floff)
+    floff, m.mlp_proj_w  = restore_floats((NBlocks,) + MlpCProjWShape , floff)
+    floff, m.mlp_proj_b  = restore_floats((NBlocks,) + MlpCProjBShape , floff)
+    floff, m.attn_w      = restore_floats((NBlocks,) + CAttnWShape    , floff)
+    floff, m.attn_b      = restore_floats((NBlocks,) + CAttnBShape    , floff)
+    floff, m.attn_proj_w = restore_floats((NBlocks,) + CProjWShape    , floff)
+    floff, m.attn_proj_b = restore_floats((NBlocks,) + CProjBShape    , floff)
+    floff, m.ln1_b       = restore_floats((NBlocks,) + BlockLnBShape  , floff)
+    floff, m.ln1_g       = restore_floats((NBlocks,) + BlockLnBShape  , floff)
+    floff, m.ln2_b       = restore_floats((NBlocks,) + BlockLnBShape  , floff)
+    floff, m.ln2_g       = restore_floats((NBlocks,) + BlockLnBShape  , floff)
+    floff, m.lnf_b       = restore_floats(ParamsLnFValShape           , floff)
+    floff, m.lnf_g       = restore_floats(ParamsLnFValShape           , floff)
+    floff, m.decoder_idx = restore_ints(DecoderIdxShape               , floff)
 
     with open("model.dat", "rb") as f:
         f.seek(floff)
         decoder_txti_ub : bytes = f.read(DecoderTxtUtf8Len)
-        decoder_txti_u = decoder_txti_ub.decode("utf-8")
-        assert len(decoder_txti_u) == DecoderTxtAsciiLen
-        assert decoder_txti_u == decoder_txt
-        m.decoder_idx = decoder_txti_u
+        m.decoder_txt = decoder_txti_ub.decode("utf-8")
+        assert len(m.decoder_txt) == DecoderTxtAsciiLen
         floff += DecoderTxtUtf8Len
 
-    floff, vocab_idxi = restore_ints(VocabIdxShape, floff)
-    assert np.all(vocab_idxi == vocab_idx)
-    m.vocab_idx = vocab_idxi
+    floff, m.vocab_idx = restore_ints(VocabIdxShape, floff)
 
     with open("model.dat", "rb") as f:
         f.seek(floff)
         vocab_txt_ub : bytes = f.read(VocabTxtUtf8Len)
-        vocab_txt_u = vocab_txt_ub.decode("utf-8")
-        assert len(vocab_txt_u) == VocabTxtAsciiLen
-        assert vocab_txt_u == vocab_txt
-        m.vocab_txt = vocab_txt_u
+        m.vocab_txt = vocab_txt_ub.decode("utf-8")
+        assert len(m.vocab_txt) == VocabTxtAsciiLen
         floff += VocabTxtUtf8Len
 
-    floff, byte_decoderi = restore_ints(DecoderShape, floff)
-    assert np.all(byte_decoderi == byte_decoder)
-    m.byte_decoder = byte_decoderi
-
-    t2 = clock()
-    print("Restore time: ", t2 - t1)
+    floff, m.byte_decoder = restore_ints(DecoderShape, floff)
 
-    return mo
+    return m
 
 
 def prod_tuple(t : tuple[int]) -> int:
@@ -604,37 +592,31 @@ def make_empty_model_with_metadata(
     return mo
 
 
-def check_model(mo,
-                n_vocab,
-                decoder_idx,
-                vocab_idx,
-                byte_decoder,
-                n_embd,
-                nblocks):
+def check_model(mo : Model) -> None:
     check_model_metadata(mo)
 
-    assert mo.mlp_fc_w.shape    == (nblocks,) + MlpCFcWShape
-    assert mo.mlp_fc_b.shape    == (nblocks,) + MlpCFcBShape
-    assert mo.mlp_proj_w.shape  == (nblocks,) + MlpCProjWShape
-    assert mo.mlp_proj_b.shape  == (nblocks,) + MlpCProjBShape
-    assert mo.attn_w.shape      == (nblocks,) + CAttnWShape
-    assert mo.attn_b.shape      == (nblocks,) + CAttnBShape
-    assert mo.attn_proj_w.shape == (nblocks,) + CProjWShape
-    assert mo.attn_proj_b.shape == (nblocks,) + CProjBShape
-    assert mo.ln1_g.shape       == (nblocks,) + BlockLnGShape
-    assert mo.ln1_b.shape       == (nblocks,) + BlockLnBShape
-    assert mo.ln2_g.shape       == (nblocks,) + BlockLnGShape
-    assert mo.ln2_b.shape       == (nblocks,) + BlockLnBShape
-    assert mo.wte.shape         == ParamsWteShape
-    assert mo.wpe.shape         == ParamsWpeShape
-    assert mo.lnf_g.shape       == BlockLnGShape
-    assert mo.lnf_b.shape       == BlockLnBShape
-
-    assert n_vocab              == ParamsWteShape[0] == mo.n_vocab
-    assert np.size(mo.wte, 1)   == n_embd
-    assert decoder_idx.shape    == (mo.decoder_idx_len,)
-    assert vocab_idx.shape      == (mo.vocab_idx_len,)
-    assert byte_decoder.shape   == (mo.byte_decoder_len,)
+    assert mo.mlp_fc_w.shape     == (NBlocks,) + MlpCFcWShape
+    assert mo.mlp_fc_b.shape     == (NBlocks,) + MlpCFcBShape
+    assert mo.mlp_proj_w.shape   == (NBlocks,) + MlpCProjWShape
+    assert mo.mlp_proj_b.shape   == (NBlocks,) + MlpCProjBShape
+    assert mo.attn_w.shape       == (NBlocks,) + CAttnWShape
+    assert mo.attn_b.shape       == (NBlocks,) + CAttnBShape
+    assert mo.attn_proj_w.shape  == (NBlocks,) + CProjWShape
+    assert mo.attn_proj_b.shape  == (NBlocks,) + CProjBShape
+    assert mo.ln1_g.shape        == (NBlocks,) + BlockLnGShape
+    assert mo.ln1_b.shape        == (NBlocks,) + BlockLnBShape
+    assert mo.ln2_g.shape        == (NBlocks,) + BlockLnGShape
+    assert mo.ln2_b.shape        == (NBlocks,) + BlockLnBShape
+    assert mo.wte.shape          == ParamsWteShape
+    assert mo.wpe.shape          == ParamsWpeShape
+    assert mo.lnf_g.shape        == BlockLnGShape
+    assert mo.lnf_b.shape        == BlockLnBShape
+
+    assert mo.n_vocab            == ParamsWteShape[0]
+    assert np.size(mo.wte, 1)    == NEmbed
+    assert mo.decoder_idx.shape  == (mo.decoder_idx_len,)
+    assert mo.vocab_idx.shape    == (mo.vocab_idx_len,)
+    assert mo.byte_decoder.shape == (mo.byte_decoder_len,)
 
 def check_model_metadata(mo):
     assert mo.model_type       == ModelType
@@ -738,8 +720,8 @@ def main(model_size: str = "124M",
     print("Converting model, saving to `model.dat`")
     t1 = clock()
 
-    idx : np.ndarray = load_decoder_idx(decoder)
-    assert idx.shape == DecoderIdxShape
+    decoder_idx : np.ndarray = load_decoder_idx(decoder)
+    assert decoder_idx.shape == DecoderIdxShape
 
     decoder_txt : str = "".join(decoder)
     assert len(decoder_txt) == DecoderTxtAsciiLen
@@ -757,7 +739,7 @@ def main(model_size: str = "124M",
         convert(params,
                 hparams["n_head"],
                 hparams["n_ctx"],
-                idx,
+                decoder_idx,
                 decoder_txt,
                 vocab_idx,
                 vocab_txt,
diff --git a/lp-fastgpt/main.py b/lp-fastgpt/main.py
index 71245a4..87bc256 100644
--- a/lp-fastgpt/main.py
+++ b/lp-fastgpt/main.py
@@ -1,21 +1,6 @@
-import create_model
-
-import numpy as np
-
-
-def load_input(filename: str, ) -> tuple[str, int]:
-    """Straight transcription of load_input from
-    driver.f90."""
-    n_tokens_to_generate: int = 20
-    input_txt: str = '''Alan Turing theorized that
-    computers would one day become very powerful
-    , but even he could not imagine'''
-    # input_txt2: str = ''
-    # u : int = 0
-    # ios : int = 0
-    return input_txt, n_tokens_to_generate
+from restore_model import Model, restore_model
 
 
 if __name__ == '__main__':
     print("hello lp_fastgpt")
-    create_model.main()
+    m : Model = restore_model()
diff --git a/lp-fastgpt/restore_model.py b/lp-fastgpt/restore_model.py
new file mode 100644
index 0000000..3825762
--- /dev/null
+++ b/lp-fastgpt/restore_model.py
@@ -0,0 +1,402 @@
+"""This script restores the model from `model.dat`, which
+contains everything (all the parameters, all the weights,
+encoding/decoding information).
+
+Parts of this script were taken from the picoGPT project:
+https://github.com/jaymody/picoGPT
+
+Those are licensed as:
+
+MIT License
+
+Copyright (c) 2023 Jay Mody
+
+Permission is hereby granted, free of charge, to any person
+obtaining a copy of this software and associated documentation
+files (the "Software"), to deal in the Software without
+restriction, including without limitation the rights to use,
+copy, modify, merge, publish, distribute, sublicense, and/or
+sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the
+Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
+KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
+PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+"""
+
+
+from time import monotonic as clock
+import numpy as np
+from typing import Union
+
+from dataclasses import dataclass
+
+
+# === Magic (Unexplained) Numbers =========================
+
+DecoderIdxType     = np.ndarray
+DecoderIdxShape    = (50_258,)
+
+DecoderTxtType     = str
+DecoderTxtAsciiLen = 320_827
+DecoderTxtUtf8Len  = 356_735  # It's checked somewhere below.
+
+VocabIdxType       = np.ndarray
+VocabIdxShape      = (50_002,)
+
+VocabTxtType       = str
+VocabTxtAsciiLen   = 370_558
+VocabTxtUtf8Len    = 406_304  # It's checked somewhere below.
+NVocab             = 50_257
+
+DecoderShape       = (256,)
+DecoderLen         = DecoderShape[0]
+
+NBlocks            = 12
+
+ModelType          = 0xfa51697
+ModelVersion       = 1
+
+NCtx               = 1024
+NEmbed             = 768
+
+Magic2304          = 2304
+Magic3072          = 3072
+
+
+# === Block-Attention components ======================
+CAttnBType         = np.ndarray
+CAttnBShape        = (        Magic2304,)
+CAttnWType         = np.ndarray
+CAttnWShape        = (NEmbed, Magic2304,)
+CAttnType          = dict[str, Union[CAttnBType, CAttnWType]]
+
+CProjBType         = np.ndarray
+CProjBShape        = (        NEmbed,)
+CProjWType         = np.ndarray
+CProjWShape        = (NEmbed, NEmbed,)
+CProjType          = dict[str, Union[CProjBType, CProjWType]]
+# --- top level
+BlockAttnType      = dict[str, Union[CAttnType, CProjType]]
+# === BlockLn components ==============================
+BlockLnBType       = np.ndarray
+BlockLnBShape      = (NEmbed,)
+BlockLnGType       = np.ndarray
+BlockLnGShape      = (NEmbed,)
+# --- top level
+BlockLnType        = dict[str, Union[BlockLnBType, BlockLnGType]]
+# === BlockMlp components =============================
+MlpCFcBType        = np.ndarray
+MlpCFcBShape       = (        Magic3072,)
+MlpCFcWType        = np.ndarray
+MlpCFcWShape       = (NEmbed, Magic3072,)
+MlpCFcType         = dict[str, Union[MlpCFcBType, MlpCFcWType]]
+
+MlpCProjBType      = np.ndarray
+MlpCProjBShape     = (          NEmbed,)
+MlpCProjWType      = np.ndarray
+MlpCProjWShape     = (Magic3072, NEmbed,)
+MlpCProjType       = dict[str, Union[MlpCProjBType, MlpCProjWType]]
+# --- top level
+BlockMlpType       = dict[str, Union[MlpCFcType, MlpCProjType]]
+# =====================================================
+ParamsBlockType    = dict[str, Union[BlockAttnType,
+                                     BlockLnType,  # two of these
+                                     BlockMlpType]]
+
+ParamsBlocksType   = list[ParamsBlockType]
+ParamsLnFType      = dict[str, np.ndarray]
+ParamsLnFValShape  = BlockLnBShape
+ParamsWpeType      = np.ndarray
+ParamsWpeShape     = (  NCtx, NEmbed,)
+ParamsWteType      = np.ndarray
+ParamsWteShape     = (NVocab, NEmbed,)
+ParamsType         = dict[str, Union[ParamsBlocksType,
+                                     ParamsLnFType,
+                                     ParamsWpeType,
+                                     ParamsWteType]]
+
+ModelMetadataType  = np.ndarray
+ModelMetadataShape = (12,)
+
+HParamsType = dict[str, int]
+
+
+@dataclass
+class Model:
+    # integer metadata
+    model_type       : int
+    model_version    : int
+    n_vocab          : int
+    n_ctx            : int
+    n_embd           : int
+    n_layer          : int
+    n_head           : int
+    decoder_idx_len  : int
+    decoder_txt_len  : int
+    vocab_idx_len    : int
+    vocab_txt_len    : int
+    byte_decoder_len : int
+    # check shapes in asserts for now; type system too weak.
+    mlp_fc_w         : np.ndarray
+    mlp_fc_b         : np.ndarray
+    mlp_proj_w       : np.ndarray
+    mlp_proj_b       : np.ndarray
+    attn_w           : np.ndarray
+    attn_b           : np.ndarray
+    attn_proj_w      : np.ndarray
+    attn_proj_b      : np.ndarray
+    ln1_g            : np.ndarray
+    ln1_b            : np.ndarray
+    ln2_g            : np.ndarray
+    ln2_b            : np.ndarray
+    wte              : np.ndarray
+    wpe              : np.ndarray
+    lnf_g            : np.ndarray
+    lnf_b            : np.ndarray
+    # auxiliary matrices and texts
+    decoder_idx      : np.ndarray
+    decoder_txt      : str
+    vocab_idx        : np.ndarray
+    vocab_txt        : str
+    byte_decoder     : np.ndarray
+
+
+def restore_model() -> Model:
+    m : Model = make_empty_model_with_metadata(
+        model_type       = 0,
+        model_version    = 0,
+        n_vocab          = 0,
+        n_ctx            = 0,
+        n_embd           = 0,
+        n_layer          = 0,
+        n_head           = 0,
+        decoder_idx_len  = 0,
+        decoder_txt_len  = 0,
+        vocab_idx_len    = 0,
+        vocab_txt_len    = 0,
+        byte_decoder_len = 0,
+        )
+
+    floff : int = 0
+    metadata : np.ndarray = (
+        np.fromfile("model.dat",
+                    dtype=np.int32,
+                    count=ModelMetadataShape[0],
+                    offset=floff))
+
+    m.model_type       = metadata[ 0]
+    m.model_version    = metadata[ 1]
+    m.n_vocab          = metadata[ 2]
+    m.n_ctx            = metadata[ 3]
+    m.n_embd           = metadata[ 4]
+    m.n_layer          = metadata[ 5]
+    m.n_head           = metadata[ 6]
+    m.decoder_idx_len  = metadata[ 7]
+    m.decoder_txt_len  = metadata[ 8]
+    m.vocab_idx_len    = metadata[ 9]
+    m.vocab_txt_len    = metadata[10]
+    m.byte_decoder_len = metadata[11]
+
+    check_model_metadata(m)
+
+    floff += ModelMetadataShape[0] * BYTES_PER_INT32
+    floff, m.wte         = restore_floats(ParamsWteShape              , floff)
+    floff, m.wpe         = restore_floats(ParamsWpeShape              , floff)
+    floff, m.mlp_fc_w    = restore_floats((NBlocks,) + MlpCFcWShape   , floff)
+    floff, m.mlp_fc_b    = restore_floats((NBlocks,) + MlpCFcBShape   , floff)
+    floff, m.mlp_proj_w  = restore_floats((NBlocks,) + MlpCProjWShape , floff)
+    floff, m.mlp_proj_b  = restore_floats((NBlocks,) + MlpCProjBShape , floff)
+    floff, m.attn_w      = restore_floats((NBlocks,) + CAttnWShape    , floff)
+    floff, m.attn_b      = restore_floats((NBlocks,) + CAttnBShape    , floff)
+    floff, m.attn_proj_w = restore_floats((NBlocks,) + CProjWShape    , floff)
+    floff, m.attn_proj_b = restore_floats((NBlocks,) + CProjBShape    , floff)
+    floff, m.ln1_b       = restore_floats((NBlocks,) + BlockLnBShape  , floff)
+    floff, m.ln1_g       = restore_floats((NBlocks,) + BlockLnBShape  , floff)
+    floff, m.ln2_b       = restore_floats((NBlocks,) + BlockLnBShape  , floff)
+    floff, m.ln2_g       = restore_floats((NBlocks,) + BlockLnBShape  , floff)
+    floff, m.lnf_b       = restore_floats(ParamsLnFValShape           , floff)
+    floff, m.lnf_g       = restore_floats(ParamsLnFValShape           , floff)
+    floff, m.decoder_idx = restore_ints(DecoderIdxShape               , floff)
+
+    with open("model.dat", "rb") as f:
+        f.seek(floff)
+        decoder_txti_ub : bytes = f.read(DecoderTxtUtf8Len)
+        m.decoder_txt = decoder_txti_ub.decode("utf-8")
+        assert len(m.decoder_txt) == DecoderTxtAsciiLen
+        floff += DecoderTxtUtf8Len
+
+    floff, m.vocab_idx = restore_ints(VocabIdxShape, floff)
+
+    with open("model.dat", "rb") as f:
+        f.seek(floff)
+        vocab_txt_ub : bytes = f.read(VocabTxtUtf8Len)
+        m.vocab_txt = vocab_txt_ub.decode("utf-8")
+        assert len(m.vocab_txt) == VocabTxtAsciiLen
+        floff += VocabTxtUtf8Len
+
+    floff, m.byte_decoder = restore_ints(DecoderShape, floff)
+
+    return m
+
+
+def prod_tuple(t : tuple[int]) -> int:
+    result : int = 1
+    for e in t:
+        result *= e
+    return result
+
+
+BYTES_PER_INT32   = 4
+BYTES_PER_FLOAT32 = 4
+
+
+def restore_floats(shape : tuple, offset : int) -> tuple[int, np.ndarray]:
+    """agnostic to length of shape; TODO impossible to statically type"""
+    result : np.ndarray
+    count = prod_tuple(shape)
+    result = np.fromfile("model.dat",
+                         dtype=np.float32,
+                         count=count,
+                         offset=offset)
+    result = np.reshape(result, shape)
+    new_offset : int = offset + (count * BYTES_PER_FLOAT32)
+    return new_offset, result
+
+
+def restore_ints(shape : tuple, offset : int) -> tuple[int, np.ndarray]:
+    """agnostic to length of shape; TODO impossible to statically type"""
+    result : np.ndarray
+    count = prod_tuple(shape)
+    result = np.fromfile("model.dat",
+                         dtype=np.int32,
+                         count=count,
+                         offset=offset)
+    result = np.reshape(result, shape)
+    new_offset : int = offset + (count * BYTES_PER_INT32)
+    return new_offset, result
+
+
+def make_empty_model_with_metadata(
+        model_type       : int,
+        model_version    : int,
+        n_vocab          : int,
+        n_ctx            : int,
+        n_embd           : int,
+        n_layer          : int,
+        n_head           : int,
+        decoder_idx_len  : int,
+        decoder_txt_len  : int,
+        vocab_idx_len    : int,
+        vocab_txt_len    : int,
+        byte_decoder_len : int,) -> Model:
+
+    mo: Model = Model(
+        model_type       = model_type,
+        model_version    = model_version,
+        n_vocab          = n_vocab,
+        n_ctx            = n_ctx,
+        n_embd           = n_embd,
+        n_layer          = n_layer,
+        n_head           = n_head,
+        decoder_idx_len  = decoder_idx_len,
+        decoder_txt_len  = decoder_txt_len,
+        vocab_idx_len    = vocab_idx_len,
+        vocab_txt_len    = vocab_txt_len,
+        byte_decoder_len = byte_decoder_len,
+
+        mlp_fc_w     = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32),
+        mlp_fc_b     = np.empty((n_layer, 4 * n_embd)         , dtype=np.float32),
+        mlp_proj_w   = np.empty((n_layer, 4 * n_embd, n_embd) , dtype=np.float32),
+        mlp_proj_b   = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        attn_w       = np.empty((n_layer, n_embd, 3 * n_embd) , dtype=np.float32),
+        attn_b       = np.empty((n_layer, 3 * n_embd)         , dtype=np.float32),
+        attn_proj_w  = np.empty((n_layer, n_embd, n_embd)     , dtype=np.float32),
+        attn_proj_b  = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln1_g        = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln1_b        = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln2_g        = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        ln2_b        = np.empty((n_layer, n_embd)             , dtype=np.float32),
+        wte          = np.empty(0                             , dtype=np.float32),
+        wpe          = np.empty(0                             , dtype=np.float32),
+        lnf_g        = np.empty(0                             , dtype=np.float32),
+        lnf_b        = np.empty(0                             , dtype=np.float32),
+
+        decoder_idx  = np.empty(0                             , dtype=np.int32),
+        decoder_txt  = '',
+        vocab_idx    = np.empty(0                             , dtype=np.int32),
+        vocab_txt    = '',
+        byte_decoder = np.empty(0                             , dtype=np.int32),
+    )
+    return mo
+
+
+def check_model(mo : Model) -> None:
+    check_model_metadata(mo)
+
+    assert mo.mlp_fc_w.shape     == (NBlocks,) + MlpCFcWShape
+    assert mo.mlp_fc_b.shape     == (NBlocks,) + MlpCFcBShape
+    assert mo.mlp_proj_w.shape   == (NBlocks,) + MlpCProjWShape
+    assert mo.mlp_proj_b.shape   == (NBlocks,) + MlpCProjBShape
+    assert mo.attn_w.shape       == (NBlocks,) + CAttnWShape
+    assert mo.attn_b.shape       == (NBlocks,) + CAttnBShape
+    assert mo.attn_proj_w.shape  == (NBlocks,) + CProjWShape
+    assert mo.attn_proj_b.shape  == (NBlocks,) + CProjBShape
+    assert mo.ln1_g.shape        == (NBlocks,) + BlockLnGShape
+    assert mo.ln1_b.shape        == (NBlocks,) + BlockLnBShape
+    assert mo.ln2_g.shape        == (NBlocks,) + BlockLnGShape
+    assert mo.ln2_b.shape        == (NBlocks,) + BlockLnBShape
+    assert mo.wte.shape          == ParamsWteShape
+    assert mo.wpe.shape          == ParamsWpeShape
+    assert mo.lnf_g.shape        == BlockLnGShape
+    assert mo.lnf_b.shape        == BlockLnBShape
+
+    assert mo.n_vocab            == ParamsWteShape[0]
+    assert np.size(mo.wte, 1)    == NEmbed
+    assert mo.decoder_idx.shape  == (mo.decoder_idx_len,)
+    assert mo.vocab_idx.shape    == (mo.vocab_idx_len,)
+    assert mo.byte_decoder.shape == (mo.byte_decoder_len,)
+
+
+def check_model_metadata(mo):
+    assert mo.model_type       == ModelType
+    assert mo.model_version    == ModelVersion
+    assert mo.n_vocab          == NVocab
+    assert mo.n_ctx            == NCtx
+    assert mo.n_embd           == NEmbed
+    assert mo.n_layer          == NBlocks
+    assert mo.n_head           == NBlocks
+
+    assert mo.decoder_idx_len  == DecoderIdxShape[0]
+    assert mo.decoder_txt_len  == DecoderTxtUtf8Len
+    assert mo.vocab_idx_len    == VocabIdxShape[0]
+    assert mo.vocab_txt_len    == VocabTxtUtf8Len
+    assert mo.byte_decoder_len == DecoderShape[0]
+
+
+def main() -> Model:
+
+    print("Restoring model")
+    t1 = clock()
+    m : Model = restore_model()
+    t2 = clock()
+    print("  Done. Time: ", t2 - t1)
+    # ================================================================
+
+    return m
+
+
+if __name__ == "__main__":
+    import fire
+    fire.Fire(main)

From b77ed790b877bb4d66038f781ca2b2c4deffa829 Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Sat, 5 Aug 2023 16:55:04 -0700
Subject: [PATCH 81/84] starting the tokenizer

---
 driver.f90                  |  12 +-
 lp-fastgpt/restore_model.py |  94 +++++--
 lp-fastgpt/timer.py         |  36 +++
 tokenizer.f90               | 532 ++++++++++++++++++------------------
 4 files changed, 380 insertions(+), 294 deletions(-)
 create mode 100644 lp-fastgpt/timer.py

diff --git a/driver.f90 b/driver.f90
index c4bf9f0..9c92238 100644
--- a/driver.f90
+++ b/driver.f90
@@ -122,13 +122,23 @@ subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output)
 
         ! Compute byte_decoder:
         allocate(byte_decoder(0:maxval(m%byte_encoder)))
+        print "(a)", "size(m%byte_encoder)"
+        print *, size(m%byte_encoder)
+        print "(a)", "byte encoder max"
+        print *, maxval(m%byte_encoder)
+        print "(a)", "size(byte_decoder)"
+        print *, size(byte_decoder)
+!        print "(a)", "byte encoder"
+!        print *, m%byte_encoder
         byte_decoder = 0
         do i = 0, size(m%byte_encoder) - 1
             byte_decoder(m%byte_encoder(i)) = i
         end do
+!        print "(a)", "byte decoder"
+!        print *, byte_decoder
 
         print "(a)", "Input text"
-        print "(a)", input_txt
+        print *, input_txt
 
         print *
         print "(a)", "Encoding: tokenizing input text into tokens (currently slow)..."
diff --git a/lp-fastgpt/restore_model.py b/lp-fastgpt/restore_model.py
index 3825762..21373c7 100644
--- a/lp-fastgpt/restore_model.py
+++ b/lp-fastgpt/restore_model.py
@@ -35,16 +35,15 @@
 
 """
 
-
-from time import monotonic as clock
 import numpy as np
 from typing import Union
 
 from dataclasses import dataclass
 
+from timer import Timer
 
-# === Magic (Unexplained) Numbers =========================
 
+# === Magic (Unexplained) Numbers =========================
 DecoderIdxType     = np.ndarray
 DecoderIdxShape    = (50_258,)
 
@@ -74,7 +73,6 @@
 Magic2304          = 2304
 Magic3072          = 3072
 
-
 # === Block-Attention components ======================
 CAttnBType         = np.ndarray
 CAttnBShape        = (        Magic2304,)
@@ -89,6 +87,7 @@
 CProjType          = dict[str, Union[CProjBType, CProjWType]]
 # --- top level
 BlockAttnType      = dict[str, Union[CAttnType, CProjType]]
+
 # === BlockLn components ==============================
 BlockLnBType       = np.ndarray
 BlockLnBShape      = (NEmbed,)
@@ -96,6 +95,7 @@
 BlockLnGShape      = (NEmbed,)
 # --- top level
 BlockLnType        = dict[str, Union[BlockLnBType, BlockLnGType]]
+
 # === BlockMlp components =============================
 MlpCFcBType        = np.ndarray
 MlpCFcBShape       = (        Magic3072,)
@@ -104,12 +104,13 @@
 MlpCFcType         = dict[str, Union[MlpCFcBType, MlpCFcWType]]
 
 MlpCProjBType      = np.ndarray
-MlpCProjBShape     = (          NEmbed,)
+MlpCProjBShape     = (           NEmbed,)
 MlpCProjWType      = np.ndarray
 MlpCProjWShape     = (Magic3072, NEmbed,)
 MlpCProjType       = dict[str, Union[MlpCProjBType, MlpCProjWType]]
 # --- top level
 BlockMlpType       = dict[str, Union[MlpCFcType, MlpCProjType]]
+
 # =====================================================
 ParamsBlockType    = dict[str, Union[BlockAttnType,
                                      BlockLnType,  # two of these
@@ -130,8 +131,8 @@
 ModelMetadataType  = np.ndarray
 ModelMetadataShape = (12,)
 
-HParamsType = dict[str, int]
-
+NTokensToGenerate  =   20
+MaxTokens          = 2048
 
 @dataclass
 class Model:
@@ -147,7 +148,7 @@ class Model:
     decoder_txt_len  : int
     vocab_idx_len    : int
     vocab_txt_len    : int
-    byte_decoder_len : int
+    byte_encoder_len : int
     # check shapes in asserts for now; type system too weak.
     mlp_fc_w         : np.ndarray
     mlp_fc_b         : np.ndarray
@@ -170,11 +171,11 @@ class Model:
     decoder_txt      : str
     vocab_idx        : np.ndarray
     vocab_txt        : str
-    byte_decoder     : np.ndarray
+    byte_encoder     : np.ndarray
 
 
 def restore_model() -> Model:
-    m : Model = make_empty_model_with_metadata(
+    m : Model = empty_model_with_metadata(
         model_type       = 0,
         model_version    = 0,
         n_vocab          = 0,
@@ -186,7 +187,7 @@ def restore_model() -> Model:
         decoder_txt_len  = 0,
         vocab_idx_len    = 0,
         vocab_txt_len    = 0,
-        byte_decoder_len = 0,
+        byte_encoder_len = 0,
         )
 
     floff : int = 0
@@ -207,7 +208,7 @@ def restore_model() -> Model:
     m.decoder_txt_len  = metadata[ 8]
     m.vocab_idx_len    = metadata[ 9]
     m.vocab_txt_len    = metadata[10]
-    m.byte_decoder_len = metadata[11]
+    m.byte_encoder_len = metadata[11]
 
     check_model_metadata(m)
 
@@ -246,7 +247,7 @@ def restore_model() -> Model:
         assert len(m.vocab_txt) == VocabTxtAsciiLen
         floff += VocabTxtUtf8Len
 
-    floff, m.byte_decoder = restore_ints(DecoderShape, floff)
+    floff, m.byte_encoder = restore_ints(DecoderShape, floff)
 
     return m
 
@@ -288,7 +289,7 @@ def restore_ints(shape : tuple, offset : int) -> tuple[int, np.ndarray]:
     return new_offset, result
 
 
-def make_empty_model_with_metadata(
+def empty_model_with_metadata(
         model_type       : int,
         model_version    : int,
         n_vocab          : int,
@@ -300,7 +301,7 @@ def make_empty_model_with_metadata(
         decoder_txt_len  : int,
         vocab_idx_len    : int,
         vocab_txt_len    : int,
-        byte_decoder_len : int,) -> Model:
+        byte_encoder_len : int,) -> Model:
 
     mo: Model = Model(
         model_type       = model_type,
@@ -314,7 +315,7 @@ def make_empty_model_with_metadata(
         decoder_txt_len  = decoder_txt_len,
         vocab_idx_len    = vocab_idx_len,
         vocab_txt_len    = vocab_txt_len,
-        byte_decoder_len = byte_decoder_len,
+        byte_encoder_len = byte_encoder_len,
 
         mlp_fc_w     = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32),
         mlp_fc_b     = np.empty((n_layer, 4 * n_embd)         , dtype=np.float32),
@@ -337,7 +338,7 @@ def make_empty_model_with_metadata(
         decoder_txt  = '',
         vocab_idx    = np.empty(0                             , dtype=np.int32),
         vocab_txt    = '',
-        byte_decoder = np.empty(0                             , dtype=np.int32),
+        byte_encoder = np.empty(0                             , dtype=np.int32),
     )
     return mo
 
@@ -366,7 +367,7 @@ def check_model(mo : Model) -> None:
     assert np.size(mo.wte, 1)    == NEmbed
     assert mo.decoder_idx.shape  == (mo.decoder_idx_len,)
     assert mo.vocab_idx.shape    == (mo.vocab_idx_len,)
-    assert mo.byte_decoder.shape == (mo.byte_decoder_len,)
+    assert mo.byte_encoder.shape == (mo.byte_encoder_len,)
 
 
 def check_model_metadata(mo):
@@ -382,21 +383,60 @@ def check_model_metadata(mo):
     assert mo.decoder_txt_len  == DecoderTxtUtf8Len
     assert mo.vocab_idx_len    == VocabIdxShape[0]
     assert mo.vocab_txt_len    == VocabTxtUtf8Len
-    assert mo.byte_decoder_len == DecoderShape[0]
+    assert mo.byte_encoder_len == DecoderShape[0]
+
+
+def next_token(input_ : str, i : int) -> tuple[str, int]:
+    result = ('', i)
+    if i >= len(input_):
+        return '', i
+    # elif
+    return result
+
+
+def encode(m : Model, input_ : str, byte_decoder : np.ndarray) -> np.ndarray:
+    """Compare to the fortran function in tokenizer.f90."""
+    # reshape this later after counting tokens
+    tokens2  : np.ndarray = np.zeros(MaxTokens, dtype=np.int32)
+    n_tokens : int        = 0
+    i        : int        = 0  # fortran counts from 1
+    # Python does not have \p for punctuation.
+    # rex = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
+    # rex.match(input_)
+    # while True:
+    #     tmp : str = next_token(input_, i)
+    #     if tmp == '':
+    #         break
+    return tokens2
+
+
+def gpt2_driver(m : Model, input_ : str) -> str:
+    print(f"Input to the model =\n{input_}")
+    n_tokens_to_generate : int = NTokensToGenerate
+    n_seq                : int = len(input_)
+    byte_encoder_max     : int = np.max(m.byte_encoder)
+    byte_decoder : np.ndarray = \
+        np.zeros((byte_encoder_max + 1,), dtype=np.int32)
+    for i, e in enumerate(m.byte_encoder):
+        byte_decoder[e] = i
+    # check against fortran:
+    # print(f'byte_decoder = \n{byte_decoder}')
+    encoded : np.ndarray = encode(m, input_, byte_decoder)
+    result = ''
+    return result
 
 
 def main() -> Model:
 
-    print("Restoring model")
-    t1 = clock()
-    m : Model = restore_model()
-    t2 = clock()
-    print("  Done. Time: ", t2 - t1)
-    # ================================================================
+    with Timer(text="Restored the model in {:0.4f} seconds."):
+        m : Model = restore_model()
+
+    input : str = """Alan Turing theorized that computers would one day become very powerful, but even he could not imagine"""
+    with Timer(text="Ran the model in {:0.6f} seconds."):
+        result : str = gpt2_driver(m, input)
 
     return m
 
 
 if __name__ == "__main__":
-    import fire
-    fire.Fire(main)
+    main()
diff --git a/lp-fastgpt/timer.py b/lp-fastgpt/timer.py
new file mode 100644
index 0000000..123c3f7
--- /dev/null
+++ b/lp-fastgpt/timer.py
@@ -0,0 +1,36 @@
+"""See https://realpython.com/python-timer/"""
+
+
+from dataclasses import dataclass, field
+import time
+from typing import Callable, Optional
+
+
+class TimerError(Exception):
+    """custom Exception for Timer errors"""
+
+
+@dataclass
+class Timer:
+    """nestable"""
+    text        : str = "Elapsed time: {:0.4f} seconds"
+    logger      : Callable[[str], None] = print
+    _start_time : Optional[float] = \
+        field(default=None, init=False, repr=False)
+
+    def start(self) -> None:
+        self._start_time = time.perf_counter()
+
+    def stop(self) -> float:
+        elapsed_time = time.perf_counter() - self._start_time
+        self._start_time = None
+        self.logger(self.text.format(elapsed_time))
+        return elapsed_time
+
+    # Protocol methods for context manager:
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, *exc_info):
+        self.stop()
diff --git a/tokenizer.f90 b/tokenizer.f90
index f546c20..418752a 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -1,284 +1,284 @@
 module tokenizer
-implicit none
+    implicit none
 
-type :: string
-    character(:), allocatable :: s
-end type
+    type :: string
+        character(:), allocatable :: s
+    end type
 
 contains
 
-function c2s(x) result(y)
-integer(1), intent(in) :: x(:)
-integer(1) :: xx(size(x))
-character(:), allocatable :: y
-integer :: i
-xx = x
-allocate(character(size(x)) :: y)
-do i = 1, size(x)
-    y(i:i) = char(int(xx(i),4))
-end do
-end function
+    function c2s(x) result(y)
+        integer(1), intent(in) :: x(:)
+        integer(1) :: xx(size(x))
+        character(:), allocatable :: y
+        integer :: i
+        xx = x
+        allocate(character(size(x)) :: y)
+        do i = 1, size(x)
+            y(i:i) = char(int(xx(i), 4))
+        end do
+    end function
 
-function next_token(input, i) result(y)
-! TODO: tokenize exactly according to this regex:
-! re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
-! Right now we are more greedy, but the bpe() tokenizer seems to still return
-! exactly the same tokens for most inputs (it is not clear if for all inputs).
-character(*), intent(in) :: input
-integer, intent(inout) :: i
-character(:), allocatable :: y
-if (i > len(input)) then
-    y = ""
-else if (input(i:i) == " ") then
-    y = tokenize_word(input, i)
-else if (input(i:i) == "," .or. input(i:i) == ".") then
-    y = input(i:i)
-    i = i + 1
-else
-    y = tokenize_word(input, i)
-end if
-end function
+    function next_token(input, i) result(y)
+        ! TODO: tokenize exactly according to this regex:
+        ! re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
+        ! Right now we are more greedy, but the bpe() tokenizer seems to still return
+        ! exactly the same tokens for most inputs (it is not clear if for all inputs).
+        character(*), intent(in) :: input
+        integer, intent(inout) :: i
+        character(:), allocatable :: y
+        if (i > len(input)) then
+            y = ""
+        else if (input(i:i) == " ") then
+            y = tokenize_word(input, i)
+        else if (input(i:i) == "," .or. input(i:i) == ".") then
+            y = input(i:i)
+            i = i + 1
+        else
+            y = tokenize_word(input, i)
+        end if
+    end function
+
+    function tokenize_word(input, i) result(y)
+        character(*), intent(in) :: input
+        integer, intent(inout) :: i
+        character(:), allocatable :: y
+        integer :: i0
+        i0 = i
+        if (input(i:i) == " ") then
+            i = i + 1
+        end if
+        do
+            if (i > len(input)) then
+                y = input(i0:i - 1)
+                exit
+            end if
+            if (input(i:i) == " " .or. input(i:i) == "," .or. input(i:i) == ".") then
+                y = input(i0:i - 1)
+                exit
+            end if
+            i = i + 1
+        end do
+    end function
 
-function tokenize_word(input, i) result(y)
-character(*), intent(in) :: input
-integer, intent(inout) :: i
-character(:), allocatable :: y
-integer :: i0
-i0 = i
-if (input(i:i) == " ") then
-    i = i + 1
-end if
-do
-    if (i > len(input)) then
-        y = input(i0:i-1)
-        exit
-    end if
-    if (input(i:i) == " " .or. input(i:i) == "," .or. input(i:i) == ".") then
-        y = input(i0:i-1)
-        exit
-    end if
-    i = i + 1
-end do
-end function
+    function word_idx(word, idx, decoder_txt) result(token)
+        character(*), intent(in) :: word
+        integer, intent(in) :: idx(0:)
+        integer(1), intent(in) :: decoder_txt(:)
+        integer :: token
+        integer :: i
+        ! This is O(n) search instead of O(1) lookup in a dictionary, so it is slow
+        do i = 0, ubound(idx, 1) - 1
+            if (c2s(decoder_txt(idx(i) + 1:idx(i + 1))) == word) then
+                token = i
+                return
+            end if
+        end do
+        token = -1
+    end function
 
-function word_idx(word, idx, decoder_txt) result(token)
-character(*), intent(in) :: word
-integer, intent(in) :: idx(0:)
-integer(1), intent(in) :: decoder_txt(:)
-integer :: token
-integer :: i
-! This is O(n) search instead of O(1) lookup in a dictionary, so it is slow
-do i = 0, ubound(idx,1)-1
-    if (c2s(decoder_txt(idx(i)+1:idx(i+1))) == word) then
-        token = i
-        return
-    end if
-end do
-token = -1
-end function
+    subroutine codepoint_to_utf8(s, c)
+        ! UTF-32 -> UTF-8
+        character(:), allocatable, intent(inout) :: s
+        integer, intent(in) :: c
+        integer :: d1, d2
+        if (c < 128) then
+            s = s // achar(c)
+        else if (c < 2048) then
+            d1 = ior(ishft(c, -6), 192)
+            d2 = iand(ior(c, 128), 191)
+            s = s // achar(d1) // achar(d2)
+        else
+            error stop "UTF-32 range not supported"
+        end if
+    end subroutine
 
-subroutine codepoint_to_utf8(s, c)
-! UTF-32 -> UTF-8
-character(:), allocatable, intent(inout) :: s
-integer, intent(in) :: c
-integer :: d1, d2
-if (c < 128) then
-    s = s // achar(c)
-else if (c < 2048) then
-    d1 = ior(ishft(c, -6), 192)
-    d2 = iand(ior(c, 128), 191)
-    s = s // achar(d1) // achar(d2)
-else
-    error stop "UTF-32 range not supported"
-end if
-end subroutine
+    function utf8_to_codepoint(s, i) result(c)
+        ! UTF-8 -> UTF-32
+        character(*), intent(in) :: s
+        integer, intent(inout) :: i
+        integer :: c, d
+        c = iachar(s(i:i))
+        if (c >= 128) then
+            i = i + 1
+            d = iachar(s(i:i))
+            c = ior(ishft(iand(c, 31), 6), iand(d, 63))
+        end if
+        if (c >= 2048) then
+            error stop "UTF-8 range not supported"
+        end if
+    end function
 
-function utf8_to_codepoint(s, i) result(c)
-! UTF-8 -> UTF-32
-character(*), intent(in) :: s
-integer, intent(inout) :: i
-integer :: c, d
-c = iachar(s(i:i))
-if (c >= 128) then
-    i = i + 1
-    d = iachar(s(i:i))
-    c = ior(ishft(iand(c, 31), 6), iand(d, 63))
-end if
-if (c >= 2048) then
-    error stop "UTF-8 range not supported"
-end if
-end function
+    function merge_pair(intokens, idx) result(tokens)
+        ! Merge the pair `idx`
+        type(string), intent(in) :: intokens(:)
+        integer, intent(in) :: idx
+        type(string), allocatable :: tokens(:)
+        type(string) :: merged_token
+        integer :: i
+        merged_token%s = intokens(idx)%s // intokens(idx + 1)%s
+        allocate(tokens(size(intokens) - 1))
+        do i = 1, idx - 1
+            tokens(i) = intokens(i)
+        end do
+        tokens(idx) = merged_token
+        do i = idx + 2, size(intokens)
+            tokens(i - 1) = intokens(i)
+        end do
+    end function
 
-function merge_pair(intokens, idx) result(tokens)
-! Merge the pair `idx`
-type(string), intent(in) :: intokens(:)
-integer, intent(in) :: idx
-type(string), allocatable :: tokens(:)
-type(string) :: merged_token
-integer :: i
-merged_token%s = intokens(idx)%s // intokens(idx+1)%s
-allocate(tokens(size(intokens)-1))
-do i = 1, idx-1
-    tokens(i) = intokens(i)
-end do
-tokens(idx) = merged_token
-do i = idx+2, size(intokens)
-    tokens(i-1) = intokens(i)
-end do
-end function
+    function merge_utf8_pairs(intokens) result(tokens)
+        ! Merge all UTF-8 character pairs
+        type(string), intent(in) :: intokens(:)
+        type(string), allocatable :: tokens(:)
+        integer :: i, j, ic
+        logical :: one_more_pass
+        allocate(tokens(size(intokens)))
+        tokens = intokens
+        one_more_pass = .true.
+        !print *, "merge_utf8_pairs:", size(tokens)
+        !print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens))
+        j = 1
+        do while(one_more_pass)
+            one_more_pass = .false.
+            do i = j, size(tokens) - 1
+                if (len(tokens(i)%s) == 1) then
+                    ic = iachar(tokens(i)%s(1:1))
+                    if (ic < 0) ic = ic + 256
+                    if (ic >= 128) then
+                        tokens = merge_pair(tokens, i)
+                        one_more_pass = .true.
+                        j = i + 1
+                        exit
+                    end if
+                end if
+            end do
+        end do
+        !print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens))
+    end function
 
-function merge_utf8_pairs(intokens) result(tokens)
-! Merge all UTF-8 character pairs
-type(string), intent(in) :: intokens(:)
-type(string), allocatable :: tokens(:)
-integer :: i, j, ic
-logical :: one_more_pass
-allocate(tokens(size(intokens)))
-tokens = intokens
-one_more_pass = .true.
-!print *, "merge_utf8_pairs:", size(tokens)
-!print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens))
-j = 1
-do while(one_more_pass)
-    one_more_pass = .false.
-    do i = j, size(tokens)-1
-        if (len(tokens(i)%s) == 1) then
-            ic = iachar(tokens(i)%s(1:1))
-            if (ic < 0) ic = ic + 256
-            if (ic >= 128) then
-                tokens = merge_pair(tokens, i)
-                one_more_pass = .true.
-                j = i + 1
+    function bpe(token, vocab_idx, vocab_txt) result(tokens)
+        ! Takes a token as a string, and returns bpe tokens as an array of strings
+        character(*), intent(in) :: token
+        integer, intent(in) :: vocab_idx(0:)
+        integer(1), intent(in) :: vocab_txt(:)
+        type(string), allocatable :: tokens(:)
+        integer, allocatable :: pair_scores(:)
+        integer :: not_found, merge_pair_idx
+        integer :: i
+        not_found = size(vocab_idx) + 10
+        allocate(tokens(len(token)))
+        do i = 1, len(token)
+            tokens(i)%s = token(i:i)
+        end do
+        tokens = merge_utf8_pairs(tokens)
+        do
+            !print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens))
+            if (size(tokens) == 1) then
+                ! The token pairs were either all merged into one word, or the input
+                ! token was a one character word, either way we are done:
                 exit
             end if
-        end if
-    end do
-end do
-!print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens))
-end function
-
-function bpe(token, vocab_idx, vocab_txt) result(tokens)
-! Takes a token as a string, and returns bpe tokens as an array of strings
-character(*), intent(in) :: token
-integer, intent(in) :: vocab_idx(0:)
-integer(1), intent(in) :: vocab_txt(:)
-type(string), allocatable :: tokens(:)
-integer, allocatable :: pair_scores(:)
-integer :: not_found, merge_pair_idx
-integer :: i
-not_found = size(vocab_idx) + 10
-allocate(tokens(len(token)))
-do i = 1, len(token)
-    tokens(i)%s = token(i:i)
-end do
-tokens = merge_utf8_pairs(tokens)
-do
-    !print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens))
-    if (size(tokens) == 1) then
-        ! The token pairs were either all merged into one word, or the input
-        ! token was a one character word, either way we are done:
-        exit
-    end if
-    allocate(pair_scores(size(tokens)-1))
-    ! Loop over pairs
-    do i = 1, size(tokens)-1
-        pair_scores(i) = word_idx(tokens(i)%s // " " // tokens(i+1)%s, vocab_idx, vocab_txt)
-        if (pair_scores(i) == -1) pair_scores(i) = not_found
-    end do
-    merge_pair_idx = minloc(pair_scores, 1)
-    if (pair_scores(merge_pair_idx) == not_found) then
-        ! No token pair can be merged, so we are done:
-        exit
-    end if
-    !print *, pair_scores
-    !print *, merge_pair_idx, pair_scores(merge_pair_idx)
-    tokens = merge_pair(tokens, merge_pair_idx)
-    deallocate(pair_scores)
-end do
-!print *, "final tokens = ", (tokens(i)%s // " ", i=1,size(tokens))
-end function
+            allocate(pair_scores(size(tokens) - 1))
+            ! Loop over pairs
+            do i = 1, size(tokens) - 1
+                pair_scores(i) = word_idx(tokens(i)%s // " " // tokens(i + 1)%s, vocab_idx, vocab_txt)
+                if (pair_scores(i) == -1) pair_scores(i) = not_found
+            end do
+            merge_pair_idx = minloc(pair_scores, 1)
+            if (pair_scores(merge_pair_idx) == not_found) then
+                ! No token pair can be merged, so we are done:
+                exit
+            end if
+            !print *, pair_scores
+            !print *, merge_pair_idx, pair_scores(merge_pair_idx)
+            tokens = merge_pair(tokens, merge_pair_idx)
+            deallocate(pair_scores)
+        end do
+        !print *, "final tokens = ", (tokens(i)%s // " ", i=1,size(tokens))
+    end function
 
-function encode(input, idx, decoder_txt, vocab_idx, vocab_txt, byte_encoder) &
-        result(tokens2)
-character(*), intent(in) :: input
-integer, intent(in) :: idx(0:), vocab_idx(0:), byte_encoder(0:)
-integer(1), intent(in) :: decoder_txt(:), vocab_txt(:)
-integer, parameter :: max_tokens = 2048
-integer :: tokens(max_tokens)
-integer, allocatable :: tokens2(:)
-character(:), allocatable :: tmp, tmp2
-type(string), allocatable :: bpe_tokens(:)
-integer :: i, j, c, n_tokens
-n_tokens = 0
-i = 1
-do
-    tmp = next_token(input, i)
-    if (tmp == "") exit
-    tmp2 = ""
-    do j = 1, len(tmp)
-        c = iachar(tmp(j:j))
-        c = byte_encoder(c)
-        ! c is UTF-32 (4 bytes), but only the range [0, 324] is used
-        ! Encode c from UTF-32 to UTF-8. Due to the limited range
-        ! either one or two bytes of UTF-8 are appended to tmp2:
-        call codepoint_to_utf8(tmp2, c)
-    end do
-    bpe_tokens = bpe(tmp2, vocab_idx, vocab_txt)
-    do j = 1, size(bpe_tokens)
-        n_tokens = n_tokens + 1
-        if (n_tokens > max_tokens) error stop "exceeded max_tokens"
-        tokens(n_tokens) = word_idx(bpe_tokens(j)%s, idx, decoder_txt)
-    end do
-    deallocate(tmp2)
-end do
-allocate(tokens2(n_tokens))
-do i = 1, n_tokens
-    tokens2(i) = tokens(i)
-end do
-end function
+    function encode(input, idx, decoder_txt, vocab_idx, vocab_txt, byte_encoder) &
+            result(tokens2)
+        character(*), intent(in) :: input
+        integer, intent(in) :: idx(0:), vocab_idx(0:), byte_encoder(0:)
+        integer(1), intent(in) :: decoder_txt(:), vocab_txt(:)
+        integer, parameter :: max_tokens = 2048
+        integer :: tokens(max_tokens)
+        integer, allocatable :: tokens2(:)
+        character(:), allocatable :: tmp, tmp2
+        type(string), allocatable :: bpe_tokens(:)
+        integer :: i, j, c, n_tokens
+        n_tokens = 0
+        i = 1
+        do
+            tmp = next_token(input, i)
+            if (tmp == "") exit
+            tmp2 = ""
+            do j = 1, len(tmp)
+                c = iachar(tmp(j:j))
+                c = byte_encoder(c)
+                ! c is UTF-32 (4 bytes), but only the range [0, 324] is used
+                ! Encode c from UTF-32 to UTF-8. Due to the limited range
+                ! either one or two bytes of UTF-8 are appended to tmp2:
+                call codepoint_to_utf8(tmp2, c)
+            end do
+            bpe_tokens = bpe(tmp2, vocab_idx, vocab_txt)
+            do j = 1, size(bpe_tokens)
+                n_tokens = n_tokens + 1
+                if (n_tokens > max_tokens) error stop "exceeded max_tokens"
+                tokens(n_tokens) = word_idx(bpe_tokens(j)%s, idx, decoder_txt)
+            end do
+            deallocate(tmp2)
+        end do
+        allocate(tokens2(n_tokens))
+        do i = 1, n_tokens
+            tokens2(i) = tokens(i)
+        end do
+    end function
 
-function decode(tokens, idx, decoder_txt, byte_decoder) result(output)
-integer, intent(in) :: tokens(:), idx(0:), byte_decoder(0:)
-integer(1), intent(in) :: decoder_txt(:)
-character(:), allocatable :: output
-character(:), allocatable :: output2, tmp
-integer, parameter :: max_len = 4096
-integer(1) :: output3(max_len)
-integer(1), allocatable :: output4(:)
-integer :: i, j, c, pos
-pos = 0
-do i = 1, size(tokens)
-    if (tokens(i) < 0) error stop "tokens(i) < 0"
-    ! output2 = output2 // (decoder_txt(idx(tokens(i))+1:idx(tokens(i)+1)))
-    do j = idx(tokens(i))+1, idx(tokens(i)+1)
-        pos = pos + 1
-        output3(pos) = decoder_txt(j)
-    end do
-end do
-allocate(character(0) :: output2) ! Fix GFortran warning
-allocate(output4(pos))
-do j = 1, pos
-    output4(j) = output3(j)
-end do
-output2 = c2s(output4)
-i = 1
-output = ""
-do
-    ! Decode UTF-8 (one or more bytes) to UTF-32 code point (always 4 bytes),
-    ! However for GPT-2 it seems only range 0-323 is used from UTF-32.
-    c = utf8_to_codepoint(output2, i)
-    ! [0,324] -> [0,255]
-    if (c < 0 .or. c > ubound(byte_decoder,1)) then
-        print *, "Codepoint out of range for byte decoder:", c, ubound(byte_decoder,1)
-        ! We skip this, to allow LFortran to run
-        error stop "Codepoint out of range for byte decoder"
-    else
-        tmp = achar(byte_decoder(c))
-        output = output // tmp
-    end if
-    if (i == len(output2)) exit
-    i = i + 1
-end do
-end function
+    function decode(tokens, idx, decoder_txt, byte_decoder) result(output)
+        integer, intent(in) :: tokens(:), idx(0:), byte_decoder(0:)
+        integer(1), intent(in) :: decoder_txt(:)
+        character(:), allocatable :: output
+        character(:), allocatable :: output2, tmp
+        integer, parameter :: max_len = 4096
+        integer(1) :: output3(max_len)
+        integer(1), allocatable :: output4(:)
+        integer :: i, j, c, pos
+        pos = 0
+        do i = 1, size(tokens)
+            if (tokens(i) < 0) error stop "tokens(i) < 0"
+            ! output2 = output2 // (decoder_txt(idx(tokens(i))+1:idx(tokens(i)+1)))
+            do j = idx(tokens(i)) + 1, idx(tokens(i) + 1)
+                pos = pos + 1
+                output3(pos) = decoder_txt(j)
+            end do
+        end do
+        allocate(character(0) :: output2) ! Fix GFortran warning
+        allocate(output4(pos))
+        do j = 1, pos
+            output4(j) = output3(j)
+        end do
+        output2 = c2s(output4)
+        i = 1
+        output = ""
+        do
+            ! Decode UTF-8 (one or more bytes) to UTF-32 code point (always 4 bytes),
+            ! However for GPT-2 it seems only range 0-323 is used from UTF-32.
+            c = utf8_to_codepoint(output2, i)
+            ! [0,324] -> [0,255]
+            if (c < 0 .or. c > ubound(byte_decoder, 1)) then
+                print *, "Codepoint out of range for byte decoder:", c, ubound(byte_decoder, 1)
+                ! We skip this, to allow LFortran to run
+                error stop "Codepoint out of range for byte decoder"
+            else
+                tmp = achar(byte_decoder(c))
+                output = output // tmp
+            end if
+            if (i == len(output2)) exit
+            i = i + 1
+        end do
+    end function
 
 end module

From d75adc64c8d1d4a836cca88a462f9466873675ba Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Sun, 6 Aug 2023 12:26:25 -0700
Subject: [PATCH 82/84] interim -- broken

---
 lp-fastgpt/restore_model.py | 70 +++++++++++++++++++++++++++++++++----
 tokenizer.f90               |  6 ++++
 2 files changed, 70 insertions(+), 6 deletions(-)

diff --git a/lp-fastgpt/restore_model.py b/lp-fastgpt/restore_model.py
index 21373c7..27204b9 100644
--- a/lp-fastgpt/restore_model.py
+++ b/lp-fastgpt/restore_model.py
@@ -386,14 +386,64 @@ def check_model_metadata(mo):
     assert mo.byte_encoder_len == DecoderShape[0]
 
 
+def tokenize_word(input_ : str, i : int) -> tuple[str, int]:
+    result = ('', i)
+    i0 : int = i
+    if input_ == ' ':  # only one leading space ?
+        i += 1
+    while True:
+        ci: str = input_[i]
+        if i >= len(input_) or ci == ' ' or ci == '.' or ci == ',':
+            result = (input_[i0:i], i)
+            return result
+        i += 1
+
+
 def next_token(input_ : str, i : int) -> tuple[str, int]:
     result = ('', i)
+    ci : str = input_[i]
     if i >= len(input_):
-        return '', i
-    # elif
+        return result
+    elif ci == ' ':
+        result = tokenize_word(input_, i)
+    elif ci == ',' or ci == '.':
+        i += 1
+        result = (ci, i)
+    else:
+        result = tokenize_word(input_, i)
     return result
 
 
+def merge_utf8_pairs(token : str) -> str:
+    from copy import copy
+    tokens        : str = copy(token)
+    i             : int
+    j             : int
+    ic            : int
+    one_more_pass : bool = True
+
+    j = 0
+    while one_more_pass:
+        for i in range(j, len(tokens)):
+            pass
+        pass
+    return tokens
+
+
+def bpe(m : Model, token : str) -> str:
+    tokens         : str =  ''
+    pair_scores          =  []
+    not_found      : int =  0
+    merge_pair_idx : int =  0
+    i              : int =  0
+    MAGIC          : int = 10
+
+    not_found = len(m.vocab_idx) + MAGIC
+    tokens = merge_utf8_pairs(token)
+
+    return tokens
+
+
 def encode(m : Model, input_ : str, byte_decoder : np.ndarray) -> np.ndarray:
     """Compare to the fortran function in tokenizer.f90."""
     # reshape this later after counting tokens
@@ -403,10 +453,18 @@ def encode(m : Model, input_ : str, byte_decoder : np.ndarray) -> np.ndarray:
     # Python does not have \p for punctuation.
     # rex = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
     # rex.match(input_)
-    # while True:
-    #     tmp : str = next_token(input_, i)
-    #     if tmp == '':
-    #         break
+    while True:
+        tmp : str
+        tmp, i = next_token(input_, i)
+        t : str
+        for t in tmp:
+            c : int = ord(t)
+            e = m.byte_encoder[c]
+            tmp2 : bytes = t.encode("utf-8")
+            tmp3 : int = ord(tmp2)
+            pass
+        if tmp == '':
+            break
     return tokens2
 
 
diff --git a/tokenizer.f90 b/tokenizer.f90
index 418752a..8c705dd 100644
--- a/tokenizer.f90
+++ b/tokenizer.f90
@@ -170,6 +170,12 @@ function bpe(token, vocab_idx, vocab_txt) result(tokens)
         allocate(tokens(len(token)))
         do i = 1, len(token)
             tokens(i)%s = token(i:i)
+            print "(a)", "i"
+            print *, i
+            print "(a)", "tokens(i)%s"
+            print *, tokens(i)%s
+            print "(a)", "token(i:i)"
+            print *, token(i:i)
         end do
         tokens = merge_utf8_pairs(tokens)
         do

From a375154f149be2a5950fd0bf78ecc786a73d0fcd Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Sun, 6 Aug 2023 20:16:11 -0700
Subject: [PATCH 83/84] interim -- still broken

---
 lp-fastgpt/restore_model.py | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/lp-fastgpt/restore_model.py b/lp-fastgpt/restore_model.py
index 27204b9..7e20626 100644
--- a/lp-fastgpt/restore_model.py
+++ b/lp-fastgpt/restore_model.py
@@ -414,20 +414,38 @@ def next_token(input_ : str, i : int) -> tuple[str, int]:
     return result
 
 
+def merge_pair(intokens : bytearray, idx : int) -> bytearray:
+    tokens       : bytearray
+    merged_token : bytearray = bytearray(2)
+    i            : str
+
+    merged_token[0] = intokens[idx]
+    merged_token[1] = intokens[idx + 1]
+    tokens = intokens.copy()
+    tokens
+    return tokens
+
+
 def merge_utf8_pairs(token : str) -> str:
-    from copy import copy
-    tokens        : str = copy(token)
+    tokens        : bytearray
     i             : int
     j             : int
     ic            : int
     one_more_pass : bool = True
 
+    tokens = bytearray(token.encode("utf-8"))
     j = 0
     while one_more_pass:
+        one_more_pass = False
         for i in range(j, len(tokens)):
-            pass
-        pass
-    return tokens
+            ic : int = tokens[i]
+            if ic < 0:
+                ic += 256
+            if ic >= 128:
+                tokens = merge_pair(tokens, i)
+                one_more_pass = True
+                j = i + 1
+    return tokens.decode("utf-8")
 
 
 def bpe(m : Model, token : str) -> str:

From 5717d50228aa729c63e1303a1f9eb54d1345e1c4 Mon Sep 17 00:00:00 2001
From: rebcabin <rebcabin@users.noreply.github.com>
Date: Tue, 8 Aug 2023 16:57:25 -0700
Subject: [PATCH 84/84] interim -- still broken

---
 lp-fastgpt/restore_model.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lp-fastgpt/restore_model.py b/lp-fastgpt/restore_model.py
index 7e20626..f9721d3 100644
--- a/lp-fastgpt/restore_model.py
+++ b/lp-fastgpt/restore_model.py
@@ -480,9 +480,8 @@ def encode(m : Model, input_ : str, byte_decoder : np.ndarray) -> np.ndarray:
             e = m.byte_encoder[c]
             tmp2 : bytes = t.encode("utf-8")
             tmp3 : int = ord(tmp2)
-            pass
         if tmp == '':
-            break
+            continue
     return tokens2