fix the two actual bugs

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 })
 }