go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/experiments/gzip-classify/main.go

/*

Copyright (c) 2023 - 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 (
	"bytes"
	"cmp"
	"compress/gzip"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"

	"github.com/urfave/cli/v2"
	"gopkg.in/yaml.v3"

	"go.charczuk.com/sdk/async"
	"go.charczuk.com/sdk/cliutil"
)

func main() {
	app := &cli.App{
		Name:  "gzip-classify",
		Usage: "Generate labels for arbitrary text based on a pre-trained model.",
		Commands: []*cli.Command{
			entry,
			train,
			test,
		},
		DefaultCommand: "test",
	}
	if err := app.Run(os.Args); err != nil {
		cliutil.Fatal(err)
	}
}

var entry = &cli.Command{
	Name:    "entry",
	Aliases: []string{"e"},
	Usage:   "emit a training dataset entry",
	Flags: []cli.Flag{
		&cli.StringSliceFlag{
			Name:     "label",
			Aliases:  []string{"l"},
			Usage:    "a label",
			Required: true,
		},
		&cli.StringFlag{
			Name:     "data",
			Aliases:  []string{"d"},
			Usage:    "the data itself",
			Required: true,
		},
	},
	Action: func(ctx *cli.Context) error {
		labels := ctx.StringSlice("label")
		data, err := cliutil.FileOrStdin(ctx.String("data"))
		if err != nil {
			return err
		}
		return yaml.NewEncoder(os.Stdout).Encode(TrainingSetEntry{
			Labels: labels,
			Data:   string(data),
		})
	},
}

var train = &cli.Command{
	Name:  "train",
	Usage: "train a dataset",
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:     "input",
			Aliases:  []string{"i"},
			Usage:    "The input path in glob pattern form",
			Required: true,
		},
		&cli.StringFlag{
			Name:    "output",
			Aliases: []string{"o"},
			Usage:   "The output path",
			Value:   "data.gz",
		},
	},
	Action: func(ctx *cli.Context) error {
		inputGlob := ctx.String("input")
		output := ctx.String("output")

		var dataset []TrainingSetEntry
		var datasetMu sync.Mutex

		inputGlobMatches, err := filepath.Glob(inputGlob)
		if err != nil {
			return err
		}

		b := new(async.Batch)
		b.SetLimit(runtime.NumCPU())
		for _, inputFile := range inputGlobMatches {
			b.Go(func() error {
				var fileData TrainingSet
				f, ferr := os.Open(inputFile)
				if ferr != nil {
					fmt.Fprintf(os.Stderr, "train; could not open input file %s: %v\n", inputFile, ferr)
					return nil
				}
				defer f.Close()

				if ferr = decodeFile(inputFile, f, &fileData); ferr != nil {
					return fmt.Errorf("train; could not decode input file %s: %w", inputFile, ferr)
				}
				for _, entry := range fileData.Entries {
					compressedData := gzipCompress([]byte(entry.Data))
					entry.CompressedDataLength = len(compressedData)
					datasetMu.Lock()
					dataset = append(dataset, entry)
					datasetMu.Unlock()
				}
				return nil
			})
		}

		if err = b.Wait(); err != nil {
			return err
		}

		of, err := os.Create(output)
		if err != nil {
			return fmt.Errorf("train; could not create output file: %w", err)
		}
		defer of.Close()

		gzw := gzip.NewWriter(of)
		if err = gob.NewEncoder(gzw).Encode(TrainingSet{Entries: dataset}); err != nil {
			return fmt.Errorf("train; could not encode output file: %w", err)
		}
		if err := gzw.Flush(); err != nil {
			return fmt.Errorf("train; error flushing output file writer: %w", err)
		}
		fmt.Printf("wrote training file: %q\n", output)
		return nil
	},
}

var test = &cli.Command{
	Name:  "test",
	Usage: "test a dataset",
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:    "dataset",
			Aliases: []string{"ds"},
			Usage:   "The input dataset path.",
			Value:   "data.gz",
		},
		&cli.StringFlag{
			Name:     "file",
			Aliases:  []string{"f"},
			Usage:    "The input test data path.",
			Required: true,
		},
		&cli.IntFlag{
			Name:  "k",
			Usage: "the number of results to consider",
			Value: 2,
		},
	},
	Action: func(ctx *cli.Context) error {
		datasetPath := ctx.String("dataset")
		var training TrainingSet
		err := func() error {
			f, ferr := os.Open(datasetPath)
			if ferr != nil {
				return ferr
			}
			defer f.Close()
			gzr, ferr := gzip.NewReader(f)
			if ferr != nil {
				return ferr
			}
			if ferr := gob.NewDecoder(gzr).Decode(&training); ferr != nil {
				return ferr
			}
			return nil
		}()
		if err != nil {
			return err
		}
		contents, err := cliutil.FileOrStdin(ctx.String("file"))
		if err != nil {
			return err
		}

		contentsCleaned := strings.TrimSpace(string(contents))
		contentsCleaned = strings.ToLower(string(contentsCleaned))

		compressedContents := gzipCompress([]byte(contentsCleaned))
		cx1 := len(compressedContents)

		distanceFrom := make([]LabelsDistance, 0, len(training.Entries))
		var distanceFromMu sync.Mutex

		b := new(async.Batch)
		b.SetLimit(runtime.NumCPU())
		bp := newBufferPool()

		testValue := func(ts TrainingSetEntry) func() error {
			return func() error {
				cx2 := ts.CompressedDataLength

				buf := bp.Get()
				defer bp.Put(buf)

				buf.Write(contents)
				buf.WriteRune(' ')
				buf.Write([]byte(ts.Data))
				merged := buf.Bytes()

				mergedCompressed := gzipCompress(merged)
				cx1x2 := len(mergedCompressed)
				ncd := float64(cx1x2-min(cx1, cx2)) / float64(max(cx1, cx2))

				distanceFromMu.Lock()
				distanceFrom = append(distanceFrom, LabelsDistance{Labels: ts.Labels, Distance: ncd})
				distanceFromMu.Unlock()
				return nil
			}
		}
		for _, ts := range training.Entries {
			b.Go(testValue(ts))
		}
		if err := b.Wait(); err != nil {
			return err
		}

		k := ctx.Int("k")
		output := make(map[string][]float64)
		for _, ld := range distanceFrom[:k] {
			for _, label := range ld.Labels {
				output[label] = append(output[label], ld.Distance)
			}
		}
		var labelDistance = make([]LabelDistance, 0, len(output))
		for label, values := range output {
			labelDistance = append(labelDistance, LabelDistance{Label: label, Distance: Avg(values...)})
		}
		sort.Slice(labelDistance, func(i0, i1 int) bool {
			return labelDistance[i0].Distance < labelDistance[i1].Distance
		})
		fmt.Println(labelDistance[0].Label)
		return nil
	},
}

// LabelDistance is a helper type.
type LabelDistance struct {
	Label    string
	Distance float64
}

// LabelsDistance is a helper type.
type LabelsDistance struct {
	Labels   []string
	Distance float64
}

func Avg(values ...float64) float64 {
	var accum float64
	for _, v := range values {
		accum += v
	}
	return accum / float64(len(values))
}

func min[T cmp.Ordered](a, b T) T {
	if a < b {
		return a
	}
	return b
}

func max[T cmp.Ordered](a, b T) T {
	if a < b {
		return a
	}
	return b
}

// TrainingSet is the fully serialized training set.
type TrainingSet struct {
	Entries []TrainingSetEntry
}

// TrainingSetEntry is an entry in our model set.
type TrainingSetEntry struct {
	Labels               []string `yaml:"labels"`
	Data                 string   `yaml:"data"`
	CompressedDataLength int      `yaml:"-"`
}

func decodeFile(inputPath string, f *os.File, fileData *TrainingSet) error {
	switch filepath.Ext(inputPath) {
	case ".json", "json":
		return json.NewDecoder(f).Decode(fileData)
	case ".yaml", "yaml", ".yml", "yml":
		return yaml.NewDecoder(f).Decode(fileData)
	default:
		return fmt.Errorf("invalid path for decoding: %s", inputPath)
	}
}

func gzipCompress(data []byte) []byte {
	buf := new(bytes.Buffer)
	gzw := gzip.NewWriter(buf)
	_, _ = gzw.Write(data)
	_ = gzw.Flush()
	return buf.Bytes()
}

func newBufferPool() *BufferPool {
	return &BufferPool{
		Pool: sync.Pool{New: func() interface{} {
			b := bytes.NewBuffer(make([]byte, 512))
			b.Reset()
			return b
		}},
	}
}

// BufferPool is a sync.Pool of bytes.Buffer.
type BufferPool struct {
	sync.Pool
}

// Get returns a pooled bytes.Buffer instance.
func (p *BufferPool) Get() *bytes.Buffer {
	return p.Pool.Get().(*bytes.Buffer)
}

// Put returns the pooled instance.
func (p *BufferPool) Put(b *bytes.Buffer) {
	b.Reset()
	p.Pool.Put(b)
}