json checkpoint
This commit is contained in:
ParthSareen
parent
089bbb537d
commit
a7c8cc06da
207
sample/fast_json.go
Normal file
207
sample/fast_json.go
Normal file
@ -0,0 +1,207 @@
|
||||
package sample
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"slices"
|
||||
|
||||
"github.com/ollama/ollama/model"
|
||||
)
|
||||
|
||||
type JSONState int
|
||||
|
||||
const (
|
||||
StateStart JSONState = iota
|
||||
StateInObject
|
||||
StateInObjectKey
|
||||
StateNewline
|
||||
StateTab
|
||||
StateSpace
|
||||
StateInString
|
||||
StateInInt
|
||||
StateInFloat
|
||||
StateInBool
|
||||
StateInNull
|
||||
StateInArray
|
||||
StateInColon
|
||||
StateInComma
|
||||
StateInStringEnd
|
||||
StateInObjectKeyEnd
|
||||
StateTerminate
|
||||
StateEnd
|
||||
)
|
||||
|
||||
func (s JSONState) String() string {
|
||||
switch s {
|
||||
case StateStart:
|
||||
return "StateStart"
|
||||
case StateInObject:
|
||||
return "StateInObject"
|
||||
case StateInObjectKey:
|
||||
return "StateInObjectKey"
|
||||
case StateInString:
|
||||
return "StateInString"
|
||||
case StateNewline:
|
||||
return "StateNewline"
|
||||
case StateTab:
|
||||
return "StateTab"
|
||||
case StateSpace:
|
||||
return "StateSpace"
|
||||
case StateInInt:
|
||||
return "StateInInt"
|
||||
case StateInFloat:
|
||||
return "StateInFloat"
|
||||
case StateInColon:
|
||||
return "StateInColon"
|
||||
case StateInBool:
|
||||
return "StateInBool"
|
||||
case StateInNull:
|
||||
return "StateInNull"
|
||||
case StateInArray:
|
||||
return "StateInArray"
|
||||
case StateEnd:
|
||||
return "StateEnd"
|
||||
case StateInComma:
|
||||
return "StateInComma"
|
||||
case StateInObjectKeyEnd:
|
||||
return "StateInObjectKeyEnd"
|
||||
case StateTerminate:
|
||||
return "StateTerminate"
|
||||
case StateInStringEnd:
|
||||
return "StateInStringEnd"
|
||||
default:
|
||||
return fmt.Sprintf("Unknown state: %d", s)
|
||||
}
|
||||
}
|
||||
|
||||
type JSONSampler struct {
|
||||
curNode *Node
|
||||
proc model.TextProcessor
|
||||
stack []*Node
|
||||
}
|
||||
|
||||
func NewJSONSampler(proc model.TextProcessor) (*JSONSampler, error) {
|
||||
// fmt.Println("Creating new JSON sampler")
|
||||
startNode, err := buildStateMachine(proc)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
js := &JSONSampler{
|
||||
curNode: startNode,
|
||||
proc: proc,
|
||||
}
|
||||
|
||||
return js, nil
|
||||
}
|
||||
|
||||
func (s *JSONSampler) UpdateState(tokenSlice []int32) error {
|
||||
// fmt.Printf("Updating state with token: %v\n", tokenSlice)
|
||||
// fmt.Printf("Current state: %s\n", s.curNode.State)
|
||||
|
||||
// fmt.Println("tokenSlice", tokenSlice)
|
||||
// todo: account for strings here
|
||||
for node, edge := range s.curNode.TransitionEdges {
|
||||
for _, validToken := range edge {
|
||||
if slices.Equal(tokenSlice, validToken) {
|
||||
s.curNode = node
|
||||
// fmt.Printf("Transitioned to state: %s\n", node.State)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
for node, edge := range s.curNode.TransitionEdges {
|
||||
for _, validToken := range edge {
|
||||
if len(validToken) == 1 && validToken[0] == -1 || validToken[0] == -2 {
|
||||
s.curNode = node
|
||||
// fmt.Printf("Accepting any token, staying in state: %s\n", node.State)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
fmt.Println("invalid token ", tokenSlice)
|
||||
return errors.New("invalid token")
|
||||
}
|
||||
|
||||
func (s *JSONSampler) Sample(logits []float64) ([]float64, error) {
|
||||
fmt.Printf("Sampling in state: %s\n", s.curNode.State)
|
||||
var err error
|
||||
|
||||
switch s.curNode.State {
|
||||
case StateTerminate:
|
||||
for i := range logits {
|
||||
if s.proc.Is(uint32(i), model.SpecialEOS) {
|
||||
logits[i] = 1.0
|
||||
} else {
|
||||
logits[i] = math.NaN()
|
||||
}
|
||||
}
|
||||
return logits, nil
|
||||
|
||||
case StateInInt:
|
||||
validStates := []int32{}
|
||||
minus, err := s.proc.Encode("-")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
digits := make([][]int32, 10)
|
||||
for i := 0; i < 10; i++ {
|
||||
digits[i], err = s.proc.Encode(fmt.Sprintf("%d", i))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
// Allow "-" and digits 0-9 at start
|
||||
for i := range logits {
|
||||
for _, d := range digits {
|
||||
if len(d) == 1 && int32(i) == d[0] {
|
||||
validStates = append(validStates, int32(i))
|
||||
}
|
||||
}
|
||||
if len(minus) == 1 && int32(i) == minus[0] {
|
||||
validStates = append(validStates, int32(i))
|
||||
}
|
||||
}
|
||||
return logits, nil
|
||||
|
||||
default:
|
||||
validStates := getValidStates(s.curNode)
|
||||
logits, err = s.maskLogits(logits, validStates)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return logits, nil
|
||||
}
|
||||
}
|
||||
|
||||
func getValidStates(node *Node) []int32 {
|
||||
validStates := []int32{}
|
||||
for _, edge := range node.TransitionEdges {
|
||||
for _, token := range edge {
|
||||
validStates = append(validStates, token...)
|
||||
}
|
||||
}
|
||||
return validStates
|
||||
}
|
||||
|
||||
func (s *JSONSampler) maskLogits(logits []float64, validStates []int32) ([]float64, error) {
|
||||
// fmt.Printf("Masking logits with valid states: %v\n", validStates)
|
||||
for i := range logits {
|
||||
isValid := false
|
||||
for _, token := range validStates {
|
||||
if token == -1 {
|
||||
// fmt.Println("Found sentinel token, returning unmasked logits")
|
||||
return logits, nil
|
||||
}
|
||||
if i == int(token) {
|
||||
// fmt.Printf("Found valid token: %d\n", token)
|
||||
isValid = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !isValid {
|
||||
logits[i] = math.NaN()
|
||||
}
|
||||
}
|
||||
return logits, nil
|
||||
}
|
||||
158
sample/state_machine.go
Normal file
158
sample/state_machine.go
Normal file
@ -0,0 +1,158 @@
|
||||
package sample
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/ollama/ollama/model"
|
||||
)
|
||||
|
||||
type token []int32
|
||||
|
||||
type Node struct {
|
||||
State JSONState
|
||||
TransitionEdges map[*Node][]token
|
||||
}
|
||||
|
||||
func NewNode(state JSONState) *Node {
|
||||
return &Node{
|
||||
State: state,
|
||||
TransitionEdges: make(map[*Node][]token),
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
startToken token
|
||||
endToken token
|
||||
stringToken token
|
||||
objectKeyToken token
|
||||
tabToken token
|
||||
spaceToken token
|
||||
newlineToken token
|
||||
newlineSpace token
|
||||
commaToken token
|
||||
commaToken2 token
|
||||
commaToken3 token
|
||||
colonToken token
|
||||
colonToken2 token
|
||||
)
|
||||
|
||||
func initTokens(proc model.TextProcessor) error {
|
||||
var err error
|
||||
startToken, err = proc.Encode("{")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
endToken, err = proc.Encode("}")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
stringToken, err = proc.Encode("\"")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
objectKeyToken, err = proc.Encode("\"")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tabToken, err = proc.Encode("\t")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
spaceToken, err = proc.Encode(" ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
newlineToken, err = proc.Encode("\n")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
newlineSpace, err = proc.Encode(" \n")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// TODO: figure out how to encode colon correctly
|
||||
colonToken, err = proc.Encode("\":")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("colonToken", colonToken)
|
||||
colonToken2, err = proc.Encode(":")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
commaToken, err = proc.Encode(",")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
commaToken2, err = proc.Encode("\",")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("commaToken2", commaToken2)
|
||||
commaToken3, err = proc.Encode("\",\"")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func buildStateMachine(proc model.TextProcessor) (*Node, error) {
|
||||
if err := initTokens(proc); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
startNode := NewNode(StateStart)
|
||||
objectNode := NewNode(StateInObject)
|
||||
objectKeyNode := NewNode(StateInObjectKey)
|
||||
objectKeyEndNode := NewNode(StateInObjectKeyEnd)
|
||||
stringNode := NewNode(StateInString)
|
||||
intNode := NewNode(StateInInt)
|
||||
commaNode := NewNode(StateInComma)
|
||||
colonNode := NewNode(StateInColon)
|
||||
stringEndNode := NewNode(StateInStringEnd)
|
||||
endNode := NewNode(StateEnd)
|
||||
terminateNode := NewNode(StateTerminate)
|
||||
|
||||
sentinelToken := token([]int32{-1})
|
||||
intSentinelToken := token([]int32{-2})
|
||||
|
||||
startNode.TransitionEdges[objectNode] = []token{startToken}
|
||||
|
||||
objectNode.TransitionEdges[objectKeyNode] = []token{stringToken}
|
||||
// objectNode.TransitionEdges[objectNode] = []token{newlineToken}
|
||||
// objectNode.TransitionEdges[objectNode] = []token{spaceToken}
|
||||
|
||||
objectKeyNode.TransitionEdges[objectKeyNode] = []token{sentinelToken}
|
||||
objectKeyNode.TransitionEdges[colonNode] = []token{colonToken, colonToken2}
|
||||
// characterize end of object key
|
||||
objectKeyNode.TransitionEdges[objectKeyEndNode] = []token{stringToken}
|
||||
|
||||
objectKeyEndNode.TransitionEdges[colonNode] = []token{colonToken}
|
||||
|
||||
// objectKeyNode.TransitionEdges[intNode] = []token{sentinelToken}
|
||||
|
||||
intNode.TransitionEdges[intNode] = []token{intSentinelToken}
|
||||
intNode.TransitionEdges[commaNode] = []token{commaToken, commaToken2}
|
||||
intNode.TransitionEdges[terminateNode] = []token{endToken}
|
||||
|
||||
commaNode.TransitionEdges[objectKeyNode] = []token{newlineToken}
|
||||
|
||||
colonNode.TransitionEdges[stringNode] = []token{stringToken}
|
||||
colonNode.TransitionEdges[intNode] = []token{intSentinelToken}
|
||||
|
||||
stringNode.TransitionEdges[stringNode] = []token{sentinelToken}
|
||||
stringNode.TransitionEdges[stringEndNode] = []token{stringToken}
|
||||
// "\""," Case
|
||||
stringNode.TransitionEdges[commaNode] = []token{commaToken2}
|
||||
|
||||
// "\"",\"" Case
|
||||
stringNode.TransitionEdges[objectKeyNode] = []token{commaToken3}
|
||||
|
||||
stringEndNode.TransitionEdges[commaNode] = []token{commaToken, commaToken2}
|
||||
stringEndNode.TransitionEdges[terminateNode] = []token{endToken}
|
||||
|
||||
endNode.TransitionEdges[terminateNode] = []token{endToken}
|
||||
|
||||
terminateNode.TransitionEdges[terminateNode] = []token{}
|
||||
return startNode, nil
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user