go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/experiments/text-chain/main.go

/*

Copyright (c) 2024 - Present. Will Charczuk. All rights reserved.
Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.

*/

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"sort"
	"strings"

	"github.com/urfave/cli/v2"
	"go.charczuk.com/sdk/slant"
)

func main() {
	textchain := &cli.App{
		Name:  "text-chain",
		Usage: "text-chain trains a model from existing text and writes sentences.",
		Commands: []*cli.Command{
			train(),
			generate(),
		},
	}
	if err := textchain.Run(os.Args); err != nil {
		fmt.Fprintf(os.Stderr, "%+v\n", err)
		os.Exit(1)
	}
}

func train() *cli.Command {
	return &cli.Command{
		Name:  "train",
		Usage: "Train a model file",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:  "input-file",
				Usage: "The input file in lines format",
			},
			&cli.StringFlag{
				Name:  "output-file",
				Usage: "The output file path",
			},
		},
		Action: func(ctx *cli.Context) error {
			slant.Print(os.Stdout, "TEXT CHAIN")

			inputFile := ctx.String("input-file")
			if inputFile == "" {
				return cli.Exit("--input-file is required", 10)
			}
			outputFile := ctx.String("output-file")
			if outputFile == "" {
				return cli.Exit("--output-file is required", 10)
			}

			inputFileReader, err := os.Open(inputFile)
			if err != nil {
				return err
			}
			defer inputFileReader.Close()

			wordCounts := make(map[string]map[string]int)
			scanner := bufio.NewScanner(inputFileReader)
			for scanner.Scan() {
				line := scanner.Text()
				lineFields := strings.Fields(line)
				if len(lineFields) < 3 {
					continue
				}
				previous := null
				previous2 := null
				for x := 0; x < len(lineFields); x++ {
					word := lineFields[x]
					if word == "file:" {
						continue
					}
					addNextWord(wordCounts, previous, previous2, word)
					previous = previous2
					previous2 = word
				}

				// make sure to add the trailing word
				addNextWord(wordCounts, previous, previous2, null)
			}

			fmt.Printf("word density map has %d words\n", len(wordCounts))

			wordDensity := make(map[string]map[string]float64)
			outputFileWriter, err := os.Create(outputFile)
			if err != nil {
				return err
			}
			defer outputFileWriter.Close()

			fmt.Printf("writing word densities to: %s\n", outputFile)

			for word, counts := range wordCounts {
				var total int
				for _, count := range counts {
					total += count
				}
				wordDensity[word] = make(map[string]float64)
				for next, count := range counts {
					wordDensity[word][next] = float64(count) / float64(total)
				}
			}
			return json.NewEncoder(outputFileWriter).Encode(wordDensity)
		},
	}
}

type modelDensities map[string]map[string]float64

func generate() *cli.Command {
	return &cli.Command{
		Name:  "generate",
		Usage: "Generate a statement from a model file",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:  "model-file",
				Usage: "The model file in json formt",
			},
		},
		Action: func(ctx *cli.Context) error {
			model, err := readModelFile(ctx.String("model-file"))
			if err != nil {
				return err
			}

			token := null
			token2 := null
			next := predict(model, token, token2)
			for next != null {
				fmt.Print(next)
				token = token2
				token2 = next
				next = predict(model, token, token2)
				if next != null {
					fmt.Print(" ")
				}
			}
			fmt.Println()
			return nil
		},
	}
}

type WordDensity struct {
	Word    string
	Density float64
}

func predict(model modelDensities, token, token2 string) string {
	key := token + " " + token2
	nextCounts, ok := model[key]
	if !ok {
		return null
	}
	var wordDensities []WordDensity
	for word, density := range nextCounts {
		if word == token2 {
			continue
		}
		wordDensities = append(wordDensities, WordDensity{Word: word, Density: density})
	}

	sort.Slice(wordDensities, func(a, b int) bool {
		return wordDensities[a].Density > wordDensities[b].Density
	})

	cutoff := rand.Float64()
	var total float64
	for _, wd := range wordDensities {
		total += wd.Density
		if total > cutoff {
			return wd.Word
		}
	}
	return null
}

func readModelFile(path string) (modelDensities, error) {
	input, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer input.Close()
	wordDensity := make(modelDensities)
	if err := json.NewDecoder(input).Decode(&wordDensity); err != nil {
		return nil, err
	}
	return wordDensity, nil
}

const null = "__null__"

func addNextWord(wordCounts map[string]map[string]int, from, from2, to string) {
	key := from + " " + from2
	_, ok := wordCounts[key]
	if !ok {
		wordCounts[key] = make(map[string]int)
	}
	wordCounts[key][to]++
}