<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"hash"
	"io"
	"log"
	"sync/atomic"

	"github.com/montanaflynn/stats"
	"github.com/pkg/errors"
)

// SegmentStat holds stats for a single segment
type SegmentStat struct {
	SeqNum uint64 // First occurrence of this segment in the file
	Length int    // Length of segment
	Freq   int    // How many times this segment occurred in the file
}

// ParseStats holds stats about the parsed file
type ParseStats struct {
	Cutpoints   []uint64               // indices at which we have cutpoints
	SegLengths  []uint64               // lengths of segments (between cutpoints)
	SegHashes   map[string]SegmentStat // map[crypto hash of seg] -> SegmentStat
	BytesParsed uint64                 // number of bytes parsed

	// internal:
	segNum    uint64    // tracks the segment numbers we've issued
	segHasher hash.Hash // used to generate the crypto hash of segments
}

// NewParseStats returns an initialized ParseStats struct
func NewParseStats(hasher hash.Hash) *ParseStats {
	return &ParseStats{
		SegHashes: make(map[string]SegmentStat),
		segNum:    uint64(0),
		segHasher: hasher,
	}
}

// UpdateStats updates the ParseStats with stats for the specified chunk
func (s *ParseStats) UpdateStats(segment []byte) string {

	// Sprint'ing the hash sum causes an allocation that is unnecessary and
	// is completely avoidable.
	//segHash := fmt.Sprintf("%X", s.segHasher.Sum(segment))
	segHash := string(s.segHasher.Sum(segment))
	segStat, there := s.SegHashes[segHash]
	if there {
		segStat.Freq++
	} else {
		//segStat = SegmentStat{SeqNum: s.segNum, Length: len(segment), Freq: 1}
		segStat.Freq = 1
		segStat.Length = len(segment)
		segStat.SeqNum = atomic.AddUint64(&s.segNum, 1)
	}
	s.SegHashes[segHash] = segStat

	// Additional book keeping (TODO: allow this to be disabled)
	s.Cutpoints = append(s.Cutpoints, s.BytesParsed+uint64(len(segment)))
	s.SegLengths = append(s.SegLengths, uint64(len(segment)))
	s.BytesParsed += uint64(len(segment))

	return segHash
}

// Print prints the specified ParseStats on the given output (io.Writer)
//
func (s ParseStats) Print(out io.Writer) error {
	segLens := make([]float64, 0, len(s.SegLengths))
	for _, s := range s.SegLengths {
		segLens = append(segLens, float64(s))
	}
	log.Println("len seglens:", len(segLens))

	med, err := stats.Median(segLens)
	if err != nil {
		return errors.Wrapf(err, "Failed compute median")
	}
	max, err := stats.Max(segLens)
	if err != nil {
		return errors.Wrapf(err, "Failed to compute max")
	}
	min, err := stats.Min(segLens)
	if err != nil {
		return errors.Wrapf(err, "Failed to compute min")<commit_msg>Remove unnecessary log line while printing stats<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"hash"
	"io"
	"sync/atomic"

	"github.com/montanaflynn/stats"
	"github.com/pkg/errors"
)

// SegmentStat holds stats for a single segment
type SegmentStat struct {
	SeqNum uint64 // First occurrence of this segment in the file
	Length int    // Length of segment
	Freq   int    // How many times this segment occurred in the file
}

// ParseStats holds stats about the parsed file
type ParseStats struct {
	Cutpoints   []uint64               // indices at which we have cutpoints
	SegLengths  []uint64               // lengths of segments (between cutpoints)
	SegHashes   map[string]SegmentStat // map[crypto hash of seg] -> SegmentStat
	BytesParsed uint64                 // number of bytes parsed

	// internal:
	segNum    uint64    // tracks the segment numbers we've issued
	segHasher hash.Hash // used to generate the crypto hash of segments
}

// NewParseStats returns an initialized ParseStats struct
func NewParseStats(hasher hash.Hash) *ParseStats {
	return &ParseStats{
		SegHashes: make(map[string]SegmentStat),
		segNum:    uint64(0),
		segHasher: hasher,
	}
}

// UpdateStats updates the ParseStats with stats for the specified chunk
func (s *ParseStats) UpdateStats(segment []byte) string {

	// Sprint'ing the hash sum causes an allocation that is unnecessary and
	// is completely avoidable.
	//segHash := fmt.Sprintf("%X", s.segHasher.Sum(segment))
	segHash := string(s.segHasher.Sum(segment))
	segStat, there := s.SegHashes[segHash]
	if there {
		segStat.Freq++
	} else {
		//segStat = SegmentStat{SeqNum: s.segNum, Length: len(segment), Freq: 1}
		segStat.Freq = 1
		segStat.Length = len(segment)
		segStat.SeqNum = atomic.AddUint64(&s.segNum, 1)
	}
	s.SegHashes[segHash] = segStat

	// Additional book keeping (TODO: allow this to be disabled)
	s.Cutpoints = append(s.Cutpoints, s.BytesParsed+uint64(len(segment)))
	s.SegLengths = append(s.SegLengths, uint64(len(segment)))
	s.BytesParsed += uint64(len(segment))

	return segHash
}

// Print prints the specified ParseStats on the given output (io.Writer)
//
func (s ParseStats) Print(out io.Writer) error {
	segLens := make([]float64, 0, len(s.SegLengths))
	for _, s := range s.SegLengths {
		segLens = append(segLens, float64(s))
	}

	med, err := stats.Median(segLens)
	if err != nil {
		return errors.Wrapf(err, "Failed compute median")
	}
	max, err := stats.Max(segLens)
	if err != nil {
		return errors.Wrapf(err, "Failed to compute max")
	}
	min, err := stats.Min(segLens)
	if err != nil {
		return errors.Wrapf(err, "Failed to compute min")
{'fraction_non_alphanumeric': 0.1136838109698235, 'fraction_numerical': 0.011007781362687418, 'mean_word_length': 3.705357142857143, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3325625', 'n_tokens_mistral': 1802, 'n_tokens_neox': 1727, 'n_words': 529}