Bob.Prediction

1module Bob.Prediction
2
3import Bob.Model
4import Data.Float32Bits
5import Model.Config
6import Model.Word32
7import Std.Float
8import Std.Foundation
9import Std.List
10import Std.Natural
11import Std.Word
12
13-- Prediction is a total semantic operation over a checked 256-logit row.
14-- Equal maxima keep the earlier byte, giving deterministic greedy decoding.
15family BobPredictionSemantics : Type 0
16constructor BobPredictionSemanticsValue
17field unrestricted bobPredictionIdentity : Bytes
18field unrestricted bobPredictionContextBytes : Nat
19field unrestricted bobPredictionOutputDomain : Nat
20field unrestricted bobPredictionTieBreakIdentity : Bytes
21
22end-family
23
24family BobArgmaxState : Type 0
25constructor BobArgmaxEmpty
26constructor BobArgmaxValue
27field unrestricted bobArgmaxPosition : Nat
28field unrestricted bobArgmaxLogit : F32
29
30end-family
31
32def bobPredictionSemantics : (family BobPredictionSemantics) =
33  (record
34    BobPredictionSemantics
35    (bobPredictionIdentity = b"greedy-next-byte")
36    (bobPredictionContextBytes = bobContextLength)
37    (bobPredictionOutputDomain = bobVocabularySize)
38    (bobPredictionTieBreakIdentity = b"lowest-byte-on-equal-logit"))
39
40-- IEEE total-order key for finite binary32 values.  Negative encodings are
41-- complemented; non-negative encodings have the sign bit set.  Unsigned word
42-- comparison then follows numerical order for every admitted finite logit.
43def bobF32OrderKey =
44  (lambda unrestricted value : F32 .
45    (app
46      (lambda unrestricted bits : StdU32 .
47        (nat-eliminate
48          (lambda unrestricted negative : Nat . StdU32)
49          (stdU32Xor bits 0x80000000)
50          (lambda unrestricted predecessor : Nat .
51            (lambda unrestricted induction : StdU32 . (stdU32Not bits)))
52          (stdF32IsNegative value)))
53      (stdF32ToWord32 value)))
54
55def bobArgmaxStep =
56  (lambda unrestricted position : Nat .
57    (lambda unrestricted logit : F32 .
58      (lambda unrestricted state : (family BobArgmaxState) .
59        (eliminate
60          BobArgmaxState
61          (lambda unrestricted current : (family BobArgmaxState) . (family BobArgmaxState))
62          state
63          (branch BobArgmaxEmpty . (constructor BobArgmaxState BobArgmaxValue position logit))
64          (branch
65            BobArgmaxValue
66            bestPosition
67            bestLogit
68            .
69            (nat-eliminate
70              (lambda unrestricted currentIsNotSmaller : Nat . (family BobArgmaxState))
71              (constructor BobArgmaxState BobArgmaxValue bestPosition bestLogit)
72              (lambda unrestricted predecessor : Nat .
73                (lambda unrestricted induction : (family BobArgmaxState) .
74                  (constructor BobArgmaxState BobArgmaxValue position logit)))
75              (naturalIsZero (stdU32LessThan (bobF32OrderKey logit) (bobF32OrderKey bestLogit)))))))))
76
77def bobGreedyIndex =
78  (lambda unrestricted logits : (family StdList F32) .
79    (stdListFold
80      F32
81      (pi unrestricted position : Nat . (family BobArgmaxState))
82      (lambda unrestricted logit : F32 .
83        (lambda unrestricted rest : (pi unrestricted position : Nat . (family BobArgmaxState)) .
84          (lambda unrestricted position : Nat .
85            (bobArgmaxStep position logit (rest (succ position))))))
86      (lambda unrestricted position : Nat . (constructor BobArgmaxState BobArgmaxEmpty))
87      logits
88      0))
89
90def bobArgmaxPositionOrZero =
91  (lambda unrestricted state : (family BobArgmaxState) .
92    (eliminate
93      BobArgmaxState
94      (lambda unrestricted current : (family BobArgmaxState) . Nat)
95      state
96      (branch BobArgmaxEmpty . 0)
97      (branch BobArgmaxValue position logit . position)))
98
99def bobGreedyByte =
100  (lambda unrestricted logits : (family StdList F32) .
101    (nat-to-byte (bobArgmaxPositionOrZero (bobGreedyIndex logits))))
102
103def bobEmptyLogits : (family StdList F32) =
104  (constructor StdList StdListEmpty F32)
105
106-- Reference forward semantics: select the table row named by the current raw
107-- byte.  Shape admission below makes the fallback unreachable for a valid
108-- model value; it remains explicit so the function is total.
109def bobSelectLogitsRow =
110  (lambda unrestricted table : (family StdList (family StdList F32)) .
111    (lambda unrestricted currentByte : Byte .
112      (stdOptionValueOr
113        (family StdList F32)
114        bobEmptyLogits
115        (stdListIndex (family StdList F32) table (byte-to-nat currentByte)))))
116
117def bobTableRowsValid =
118  (lambda unrestricted table : (family StdList (family StdList F32)) .
119    (stdListFold
120      (family StdList F32)
121      Nat
122      (lambda unrestricted row : (family StdList F32) .
123        (lambda unrestricted rest : Nat .
124          (naturalAnd (naturalEqual (stdListLength F32 row) bobVocabularySize) rest)))
125      1
126      table))
127
128def bobTableShapeValid =
129  (lambda unrestricted table : (family StdList (family StdList F32)) .
130    (naturalAnd
131      (naturalEqual (stdListLength (family StdList F32) table) bobVocabularySize)
132      (bobTableRowsValid table)))