module Bob.Model

import Component.Linear
import Representation.Schema
import Std.Foundation
import Std.List
import Std.Natural

-- Bob is a byte bigram, not a small transformer: a raw current byte selects
-- one row of a learned 256 x 256 logits table for the following byte.
family BobModelSemantics : Type 0
constructor BobModelSemanticsValue
field unrestricted bobModelInputIdentity : Bytes
field unrestricted bobModelVocabularySize : Nat
field unrestricted bobModelContextLength : Nat
field unrestricted bobModelTable : (family LinearComponent)
field unrestricted bobModelParameters : (family StdList (family ParameterSchema))
field unrestricted bobModelBiasParameters : Nat
field unrestricted bobModelPersistentStateValues : Nat

end-family

def bobVocabularySize : Nat =
  256

def bobContextLength : Nat =
  1

def bobParameterCount : Nat =
  (nat-multiply bobVocabularySize bobVocabularySize)

def bobTableComponent : (family LinearComponent) =
  (constructor
    LinearComponent
    LinearComponentOf
    b"logits-table"
    bobVocabularySize
    bobVocabularySize
    (constructor StdBool StdFalse))

def bobParameterPrefix : (family StdList Bytes) =
  (constructor StdList StdListCons Bytes b"bob" (constructor StdList StdListEmpty Bytes))

-- Component.Linear derives the one weight schema from the mechanism.  The
-- false bias field is what makes this list contain exactly one parameter.
def bobParameters : (family StdList (family ParameterSchema)) =
  (linearParameters bobParameterPrefix bobTableComponent)

def bobModel : (family BobModelSemantics) =
  (record
    BobModelSemantics
    (bobModelInputIdentity = b"raw-byte-sequence")
    (bobModelVocabularySize = bobVocabularySize)
    (bobModelContextLength = bobContextLength)
    (bobModelTable = bobTableComponent)
    (bobModelParameters = bobParameters)
    (bobModelBiasParameters = 0)
    (bobModelPersistentStateValues = 0))