<commit_before>package fuzzy

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"index/suffixarray"
	"log"
	"os"
	"regexp"
	"sort"
	"strings"
	"sync"
)

const (
	SpellDepthDefault              = 2
	SpellThresholdDefault          = 5
	SuffDivergenceThresholdDefault = 100
)

type Pair struct {
	str1 string
	str2 string
}

type Method int

const (
	MethodIsWord                   Method = 0
	MethodSuggestMapsToInput              = 1
	MethodInputDeleteMapsToDict           = 2
	MethodInputDeleteMapsToSuggest        = 3
)

type Potential struct {
	Term   string // Potential term string
	Score  int    // Score
	Leven  int    // Levenstein distance from the suggestion to the input
	Method Method // How this potential was matched
}

type Counts struct {
	Corpus int `json:"corpus"`
	Query  int `json:"query"`
}

type Model struct {
	Data                    map[string]*Counts  `json:"data"`
	Maxcount                int                 `json:"maxcount"`
	Suggest                 map[string][]string `json:"suggest"`
	Depth                   int                 `json:"depth"`
	Threshold               int                 `json:"threshold"`
	UseAutocomplete         bool                `json:"autocomplete"`
	SuffDivergence          int                 `json:"-"`
	SuffDivergenceThreshold int                 `json:"suff_threshold"`
	SuffixArr               *suffixarray.Index  `json:"-"`
	SuffixArrConcat         string              `json:"-"`
	sync.RWMutex
}

// For sorting autocomplete suggestions
// to bias the most popular first
type Autos struct {
	Results []string
	Model   *Model
}

func (a Autos) Len() int      { return len(a.Results) }
func (a Autos) Swap(i, j int) { a.Results[i], a.Results[j] = a.Results[j], a.Results[i] }

func (a Autos) Less(i, j int) bool {
	icc := a.Model.Data[a.Results[i]].Corpus
	jcc := a.Model.Data[a.Results[j]].Corpus
	icq := a.Model.Data[a.Results[i]].Query
	jcq := a.Model.Data[a.Results[j]].Query
	if icq == jcq {
		if icc == jcc {
			return a.Results[i] > a.Results[j]
		}
		return icc > jcc
	}
	return icq > jcq
}

func (m Method) String() string {
	switch m {
	case MethodIsWord:
		return "Input in dictionary"
	case MethodSuggestMapsToInput:
		return "Suggest maps to input"
	case MethodInputDeleteMapsToDict:
		return "Input delete maps to dictionary"
	case MethodInputDeleteMapsToSuggest:
		return "Input delete maps to suggest key"
	}
	return "unknown"
}

func (pot *Potential) String() string {
	return fmt.Sprintf("Term: %v\n\tScore: %v\n\tLeven: %v\n\tMethod: %v\n\n", pot.Term, pot.Score, pot.Leven, pot.Method)
}

// Create and initialise a new model
func NewModel() *Model {
	model := new(Model)
	return model.Init()
}

func (model *Model) Init() *Model {
	model.Data = make(map[string]*Counts)
	model.Suggest = make(map[string][]string)
	model.Depth = SpellDepthDefault
	model.Threshold = SpellThresholdDefault // Setting this to 1 is most accurate, but "1" is 5x more memory and 30x slower processing than "4". This is a big performance tuning knob
	model.UseAutocomplete = true            // Default is to include Autocomplete
	model.updateSuffixArr()
	model.SuffDivergenceThreshold = SuffDivergenceThresholdDefault
	return model
}

// Save a spelling model to disk
func (model *Model) Save(filename string) error {
	model.RLock()
	defer model.RUnlock()
	f, err := os.Create(filename)
	if err != nil {
		log.Println("Fuzzy model:", err)
		return err
	}
	b := bufio.NewWriter(f)
	e := json.NewEncoder(b)
	defer f.Close()
	defer b.Flush()
	err = e.Encode(model)
	if err != nil {
		log.Println("Fuzzy model:", err)
	}
	return err
}

// Save a spelling model to disk, but discard all
// entries less than the threshold number of occurences
// Much smaller and all that is used when generated
// as a once off, but not useful for incremental usage
func (model *Model) SaveLight(filename string) error {
	model.Lock()
	for term, count := range model.Data {
		if count.Corpus < model.Threshold {
			delete(model.Data, term)
		}
	}
	model.Unlock()
	return model.Save(filename)
}

// Load a saved model from disk
func Load(filename string) (*Model, error) {
	model := new(Model)
	f, err := os.Open(filename)
	if err != nil {
		return model, err
	}
	defer f.Close()
	d := json.NewDecoder(f)
	err = d.Decode(model)
	if err != nil {
		if err = model.convertOldFormat(filename); err != nil {
			return model, err
		}
	}
	model.updateSuffixArr()
	return model, nil
}
<commit_msg>add FromReader and ToWriter to Model<commit_after>package fuzzy

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"index/suffixarray"
	"io"
	"log"
	"os"
	"regexp"
	"sort"
	"strings"
	"sync"
)

const (
	SpellDepthDefault              = 2
	SpellThresholdDefault          = 5
	SuffDivergenceThresholdDefault = 100
)

type Pair struct {
	str1 string
	str2 string
}

type Method int

const (
	MethodIsWord                   Method = 0
	MethodSuggestMapsToInput              = 1
	MethodInputDeleteMapsToDict           = 2
	MethodInputDeleteMapsToSuggest        = 3
)

type Potential struct {
	Term   string // Potential term string
	Score  int    // Score
	Leven  int    // Levenstein distance from the suggestion to the input
	Method Method // How this potential was matched
}

type Counts struct {
	Corpus int `json:"corpus"`
	Query  int `json:"query"`
}

type Model struct {
	Data                    map[string]*Counts  `json:"data"`
	Maxcount                int                 `json:"maxcount"`
	Suggest                 map[string][]string `json:"suggest"`
	Depth                   int                 `json:"depth"`
	Threshold               int                 `json:"threshold"`
	UseAutocomplete         bool                `json:"autocomplete"`
	SuffDivergence          int                 `json:"-"`
	SuffDivergenceThreshold int                 `json:"suff_threshold"`
	SuffixArr               *suffixarray.Index  `json:"-"`
	SuffixArrConcat         string              `json:"-"`
	sync.RWMutex
}

// For sorting autocomplete suggestions
// to bias the most popular first
type Autos struct {
	Results []string
	Model   *Model
}

func (a Autos) Len() int      { return len(a.Results) }
func (a Autos) Swap(i, j int) { a.Results[i], a.Results[j] = a.Results[j], a.Results[i] }

func (a Autos) Less(i, j int) bool {
	icc := a.Model.Data[a.Results[i]].Corpus
	jcc := a.Model.Data[a.Results[j]].Corpus
	icq := a.Model.Data[a.Results[i]].Query
	jcq := a.Model.Data[a.Results[j]].Query
	if icq == jcq {
		if icc == jcc {
			return a.Results[i] > a.Results[j]
		}
		return icc > jcc
	}
	return icq > jcq
}

func (m Method) String() string {
	switch m {
	case MethodIsWord:
		return "Input in dictionary"
	case MethodSuggestMapsToInput:
		return "Suggest maps to input"
	case MethodInputDeleteMapsToDict:
		return "Input delete maps to dictionary"
	case MethodInputDeleteMapsToSuggest:
		return "Input delete maps to suggest key"
	}
	return "unknown"
}

func (pot *Potential) String() string {
	return fmt.Sprintf("Term: %v\n\tScore: %v\n\tLeven: %v\n\tMethod: %v\n\n", pot.Term, pot.Score, pot.Leven, pot.Method)
}

// Create and initialise a new model
func NewModel() *Model {
	model := new(Model)
	return model.Init()
}

func (model *Model) Init() *Model {
	model.Data = make(map[string]*Counts)
	model.Suggest = make(map[string][]string)
	model.Depth = SpellDepthDefault
	model.Threshold = SpellThresholdDefault // Setting this to 1 is most accurate, but "1" is 5x more memory and 30x slower processing than "4". This is a big performance tuning knob
	model.UseAutocomplete = true            // Default is to include Autocomplete
	model.updateSuffixArr()
	model.SuffDivergenceThreshold = SuffDivergenceThresholdDefault
	return model
}

// WriteTo writes a model to a Writer
func (model *Model) WriteTo(w io.Writer) (int64, error) {
	model.RLock()
	defer model.RUnlock()
	b, err := json.Marshal(model)
	if err != nil {
		return 0, err
	}
	n, err := w.Write(b)
	if err != nil {
		return int64(n), err
	}
	return int64(n), nil
}

// Save a spelling model to disk
func (model *Model) Save(filename string) error {
	model.RLock()
	defer model.RUnlock()
	f, err := os.Create(filename)
	if err != nil {
		log.Println("Fuzzy model:", err)
		return err
	}
	b := bufio.NewWriter(f)
	e := json.NewEncoder(b)
	defer f.Close()
	defer b.Flush()
	err = e.Encode(model)
	if err != nil {
		log.Println("Fuzzy model:", err)
	}
	return err
}

// Save a spelling model to disk, but discard all
// entries less than the threshold number of occurences
// Much smaller and all that is used when generated
// as a once off, but not useful for incremental usage
func (model *Model) SaveLight(filename string) error {
	model.Lock()
	for term, count := range model.Data {
		if count.Corpus < model.Threshold {
			delete(model.Data, term)
		}
	}
	model.Unlock()
	return model.Save(filename)
}

// FromReader loads a model from a Reader
func FromReader(r io.Reader) (*Model, error) {
	model := new(Model)
	d := json.NewDecoder(r)
	err := d.Decode(model)
	if err != nil {
		return nil, err
	}
	model.updateSuffixArr()
	return model, nil
}

// Load a saved model from disk
func Load(filename string) (*Model, error) {
	model := new(Model)
	f, err := os.Open(filename)
	if err != nil {
		return model, err
	}
	defer f.Close()
	d := json.NewDecoder(f)
	err = d.Decode(model)
	if err != nil {
		if err = model.convertOldFormat(filename); err != nil {
			return model, err
		}
	}
	model.updateSuffixArr()
	return model, nil
}

{'fraction_non_alphanumeric': 0.11653576735300361, 'fraction_numerical': 0.004351517724474634, 'mean_word_length': 2.9311639549436794, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '543492', 'n_tokens_mistral': 3278, 'n_tokens_neox': 3141, 'n_words': 946}