<commit_before>using System.Collections.Generic;
using System.IO;
using System.Linq;
using Resin.IO;

namespace Resin
{
    public class DocumentWriter
    {
        private bool _flushing;
        private readonly string _dir;

        // docid/fields/value
        private readonly IDictionary<string, Document> _docs;

        public DocumentWriter(string dir)
        {
            _dir = dir;
            _docs = new Dictionary<string, Document>();
        }

        public void Write(Document doc)
        {
            _docs[doc.Id] = doc; // this overwrites previous doc if same docId appears twice in the session
        }

        public void Flush(DixFile dix)
        {
            if (_flushing) return;

            _flushing = true;

            var batches = _docs.IntoBatches(1000).ToList();
            foreach (var batch in batches)
            {
                var fileId = Path.GetRandomFileName();
                var fileName = Path.Combine(_dir, fileId + ".d");
                var d = new DocFile(batch.ToDictionary(x => x.Key, y => y.Value));
                d.Save(fileName);
                foreach (var docId in d.Docs)
                {
                    dix.DocIdToFileId[docId.Key] = fileId;
                }
            }
            _docs.Clear();
        }
    }
}<commit_msg>configurable doc writer<commit_after>using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading.Tasks;
using Resin.IO;

namespace Resin
{
    public class DocumentWriter
    {
        private bool _flushing;
        private readonly string _dir;
        private readonly int _batchSize;

        // docid/fields/value
        private readonly IDictionary<string, Document> _docs;

        public DocumentWriter(string dir, int batchSize = 1000)
        {
            _dir = dir;
            _batchSize = batchSize;
            _docs = new Dictionary<string, Document>();
        }

        public void Write(Document doc)
        {
            _docs[doc.Id] = doc; // this overwrites previous doc if same docId appears twice in the session
        }

        public void Flush(DixFile dix)
        {
            if (_flushing) return;

            _flushing = true;

            var files = new Dictionary<string, DocFile>();
            var batches = _docs.IntoBatches(_batchSize).ToList();
            foreach (var batch in batches)
            {
                var fileId = Path.GetRandomFileName();
                var d = new DocFile(batch.ToDictionary(x => x.Key, y => y.Value));
                foreach (var docId in d.Docs)
                {
                    dix.DocIdToFileId[docId.Key] = fileId;
                }
                files.Add(fileId, d);
            }

            //foreach (var d in files)
            Parallel.ForEach(files, d =>
            {
                var fileName = Path.Combine(_dir, d.Key + ".d");
                d.Value.Save(fileName);

            });
            _docs.Clear();
        }
    }
}
{'fraction_non_alphanumeric': 0.09127382146439318, 'fraction_numerical': 0.0026746907388833165, 'mean_word_length': 1.5313028764805414, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 13, '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': '885070', 'n_tokens_mistral': 834, 'n_tokens_neox': 795, 'n_words': 215}