fix the two actual bugs

cmd/main.go

@@ -9,7 +9,7 @@ import (
 )
 
 func main() {
-	f := "names.txt"
+	f := "input.txt"
 	if len(os.Args) > 1 {
 		f = os.Args[1]
 	}
@@ -19,5 +19,5 @@ func main() {
 		return
 	}
 	s := string(b)
-	microgopt.Run(strings.Split(s, "\n"))
+	microgopt.Run(strings.Split(s, "\n"), 10000)
 }

go.mod

@@ -1,3 +1,3 @@
 module git.yetaga.in/alazyreader/microgopt
 
-go 1.26.0
+go 1.27.0

microgopt.go

@@ -47,7 +47,11 @@ func sum(l []*value) *value {
 	return r
 }
 
-func Run(docs []string) {
+func Run(docs []string, numSteps int) {
+	if numSteps == 0 {
+		numSteps = 1000 // default number of training steps
+	}
+
 	// remove leading and trailing whitespace in documents
 	for i := range docs {
 		docs[i] = strings.TrimSpace(docs[i])
@@ -104,7 +108,6 @@ func Run(docs []string) {
 	v := make([]float64, len(params)) // second moment buffer
 
 	// Repeat in sequence
-	numSteps := 1000 // number of training steps
 	for step := range numSteps {
 		// Take single document, tokenize it, surround it with BOS special token on both sides
 		doc := docs[step%len(docs)]
@@ -157,7 +160,7 @@ func Run(docs []string) {
 				probs[i] = l.Div(&value{data: temperature})
 			}
 			probs = softMax(probs)
-			tokenId := choose(probs)
+			tokenId = choose(probs)
 			if tokenId == BOS {
 				break
 			}
@@ -210,10 +213,11 @@ func rmsNorm(x []*value) []*value {
 	}
 	ms = ms.Div(&value{data: float64(len(x))})
 	scale := ms.Add(&value{data: 1e-5}).Pow(&value{data: -0.5})
+	ret := make([]*value, len(x))
 	for i := range x {
-		x[i] = x[i].Mul(scale)
+		ret[i] = x[i].Mul(scale)
 	}
-	return x
+	return ret
 }
 
 func gpt(tokenId int, posId int, keys [][][]*value, values [][][]*value) []*value {
@@ -421,5 +425,5 @@ func choose(p []*value) int {
 	// multiply the sample with the largest CDF value; easier than normalizing to [0,1)
 	val := rand.Float64() * cdf[len(cdf)-1]
 	// Search returns the smallest index i such that cdf[i] > val
-	return sort.Search(len(cdf), func(i int) bool { return cdf[i] > val })
+	return sort.Search(len(cdf)-1, func(i int) bool { return cdf[i] > val })
 }

readme.md

@@ -8,12 +8,12 @@ To use: `go run cmd/main.go input.txt`
 
 Differences between the Go and the Python, as well as notes more generally:
 
-* The GPT is implemented as a package and, separately, as a command-line wrapper that calls it, just to keep the algorithm separate from the invocation details.
-* The Value class is more type-safe in go, using values everywhere as opposed to mingling floats and values in the localgrad tuple.
-* The Value struct has actual tests confirming the backward propagation logic.
-* When writing the Value struct and its methods, I accidentally swapped the order of the values in the `localGrads` slice in `Mul` and tore my hair out trying to figure out where the bug was. When I broke down and asked copilot to "compare these two implementations and tell me how they differ," it managed to find the error -- but also reported three non-existent differences and told me that `slices.Backward()` doesn't exist.
-* Initial pass translating the linear algebra functions has me worried that all those value structs aren't going to be very fast...
-* Had to implement weighted random choice. <https://cybernetist.com/2019/01/24/random-weighted-draws-in-go/> made that relatively straightforward; it's a neat algorithm.
+- The GPT is implemented as a package and, separately, as a command-line wrapper that calls it, just to keep the algorithm separate from the invocation details.
+- The Value class is more type-safe in go, using values everywhere as opposed to mingling floats and values in the localgrad tuple.
+- The Value struct has actual tests confirming the backward propagation logic.
+- When writing the Value struct and its methods, I accidentally swapped the order of the values in the `localGrads` slice in `Mul` and tore my hair out trying to figure out where the bug was. When I broke down and asked copilot to "compare these two implementations and tell me how they differ," it managed to find the error -- but also reported three non-existent differences and told me that `slices.Backward()` doesn't exist.
+- Initial pass translating the linear algebra functions has me worried that all those value structs aren't going to be very fast...
+- Had to implement weighted random choice. <https://cybernetist.com/2019/01/24/random-weighted-draws-in-go/> made that relatively straightforward; it's a neat algorithm.
 
 First proper run: