wip!

sample/pushdown_runner.go

@@ -0,0 +1,147 @@
+package sample
+
+import (
+	"fmt"
+	"math"
+
+	"github.com/ollama/ollama/model"
+)
+
+type PushdownSampler struct {
+	// stateful
+	curNode        *PDANode
+	proc           model.TextProcessor
+	stateToNodeMap map[JSONState]*PDANode
+	braceStack     []rune
+}
+
+func NewPushdownSampler(proc model.TextProcessor) *PushdownSampler {
+	startNode, stateToNodeMap, err := BuildGraph(proc)
+	if err != nil {
+		panic(err)
+	}
+	err = PreComputeValidStates(stateToNodeMap, proc)
+	if err != nil {
+		panic(err)
+	}
+	// for id, node := range stateToNodeMap[StateInComma].MaskTokenIDToNode {
+	// 	token, err := proc.Decode([]int32{int32(id)})
+	// 	if err != nil {
+	// 		panic(err)
+	// 	}
+	// 	fmt.Println("id", id, "node", node, "token", token)
+	// }
+	// time.Sleep(10 * time.Second)
+	return &PushdownSampler{
+		curNode:        startNode,
+		proc:           proc,
+		stateToNodeMap: stateToNodeMap,
+		braceStack:     []rune{},
+	}
+}
+
+func (s *PushdownSampler) Sample(logits []float64) ([]float64, error) {
+	fmt.Println("sample:", s.curNode.State)
+
+	switch s.curNode.State {
+	case StateInObjectEnd:
+		// force finish if no braces left
+		if len(s.braceStack) == 0 {
+			s.curNode = NewPDANode(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
+		}
+		valid, err := s.proc.Encode("}")
+		if err != nil {
+			return nil, err
+		}
+		for i := range logits {
+			for _, token := range valid {
+				if i != int(token) {
+					logits[i] = math.NaN()
+				}
+			}
+		}
+		return logits, nil
+	// return logits, nil
+	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 StateInStringEnd:
+
+	// 	return logits, nil
+	default:
+		fmt.Println("masking logits current state", s.curNode.State)
+		logits, err := s.maskLogits(logits, s.curNode)
+		if err != nil {
+			return nil, err
+		}
+		return logits, nil
+	}
+}
+
+func (s *PushdownSampler) UpdateState(tokenSlice []int32) error {
+	fmt.Println("update state", s.curNode.State)
+
+	// TODO: need to handle end states and entering object case
+	if s.curNode.State == StateInObjectEnd {
+		fmt.Println("in object end")
+		if len(s.braceStack) > 0 {
+			s.braceStack = s.braceStack[:len(s.braceStack)-1]
+			return nil
+		}
+		s.curNode = NewPDANode(StateTerminate)
+		// TODO: return here?
+	}
+	// need this cause there could be multiple transitions
+	mappedString, err := s.proc.Decode(tokenSlice)
+	if err != nil {
+		return err
+	}
+	for _, r := range mappedString {
+		if r == rune('{') {
+			s.braceStack = append(s.braceStack, r)
+		}
+		if r == rune('}') {
+			if len(s.braceStack) == 0 || s.braceStack[len(s.braceStack)-1] != rune('{') {
+				return fmt.Errorf("unmatched closing brace")
+			}
+			s.braceStack = s.braceStack[:len(s.braceStack)-1]
+		}
+	}
+	for _, tokenID := range tokenSlice {
+		// transition to the next node
+		nextNode, ok := s.curNode.MaskTokenIDToNode[tokenID]
+		if !ok {
+			return fmt.Errorf("invalid token: %q", mappedString)
+		}
+		fmt.Println("transitioning to", nextNode)
+		s.curNode = s.stateToNodeMap[nextNode]
+	}
+	return nil
+}
+
+func (s *PushdownSampler) maskLogits(logits []float64, node *PDANode) ([]float64, error) {
+	for i := range logits {
+		_, exists := node.MaskTokenIDToNode[int32(i)]
+		if !exists {
+			logits[i] = math.NaN()
+		}
+	}
+	return logits, nil
+}
+
+// TODO: add penalties for string \n stuff