twpipeline 0.1.0

data/lib/twpipeline/cli.rb

@@ -0,0 +1,144 @@
+# frozen_string_literal: true
+
+require "optparse"
+
+module TWPipeline
+  class CLI
+    COMMANDS = %w[run check stages resources].freeze
+
+    def self.call(argv) = new(argv).call
+
+    def initialize(argv)
+      @argv = argv.dup
+      @options = {}
+    end
+
+    def call
+      parser.parse!(@argv, into: parsed = {})
+      @options = parsed.transform_keys { |key| key.to_s.tr("-", "_").to_sym }
+      TWPipeline.resources = resources
+      TWPipeline.warm_up
+
+      case @argv.shift
+      in nil
+        abort(parser.to_s)
+      in "run"
+        run
+      in "check"
+        check
+      in "stages"
+        stages
+      in "resources"
+        $stdout.puts(resources.to_s)
+      in name
+        stage(name)
+      end
+
+      0
+    rescue Error, TWFilter::Error => error
+      warn("twpipeline: #{error.message}")
+      1
+    end
+
+    private
+
+    def resources
+      @resources ||= @options[:ask] ? Resources.ask : Resources.detect(
+        cores: @options[:cores],
+        memory: @options[:memory]
+      )
+    end
+
+    def parser
+      @parser ||= OptionParser.new do |parser|
+        parser.banner = <<~USAGE
+          usage: twpipeline [options] <command|stage>
+
+          commands: #{COMMANDS.join(", ")}
+          stages:   #{Stage.ordered.map(&:slug).join(", ")}
+        USAGE
+        parser.on("--cores N", Integer)
+        parser.on("--memory SIZE", String)
+        parser.on("--ask", "prompt for cores and memory")
+        parser.on("--source NAME", String)
+        parser.on("--input PATH", String)
+        parser.on("--residue PATH", String)
+        parser.on("--allowlist PATH", String)
+        parser.on("--vocabulary PATH", String)
+        parser.on("--out PATH", String)
+        parser.on("--orders LIST", Array)
+        parser.on("--unit NAME", %w[token char])
+        parser.on("--floor N", Integer)
+        parser.on("--parts FIELD", String)
+        parser.on("--once-per FIELD", String)
+        parser.on("--block N", Integer)
+        parser.on("--ubiquity N", Integer)
+        parser.on("--host-cap N", Integer)
+        parser.on("--group FIELD", String)
+        parser.on("--policy NAME", %w[corpus publishable permissive])
+        parser.on("--drop", "remove rejected records instead of flagging them")
+        parser.on("--from STAGE", String)
+        parser.on("--to STAGE", String)
+        parser.on("--quiet")
+      end
+    end
+
+    def policy = TWFilter::Policy.public_send(@options.fetch(:policy, "corpus"))
+
+    def settings = @options.slice(
+      *%i[source input residue allowlist vocabulary out orders unit floor parts once_per block group ubiquity host_cap]
+    )
+
+    def build(klass) = klass.new(policy: policy, drop: @options.fetch(:drop, false), resources: resources, **settings)
+
+    def stage(name)
+      klass = Stage.lookup(name)
+      report(klass.slug) { build(klass).call($stdin, $stdout) }
+    end
+
+    def stages
+      Stage.ordered.each { |klass| $stdout.puts(format("%-14s %s", klass.slug, klass.name)) }
+    end
+
+    def selection
+      list = Stage.ordered
+      first = @options[:from] ? Stage.lookup(@options[:from]).order : list.first.order
+      last = @options[:to] ? Stage.lookup(@options[:to]).order : list.last.order
+      list.select { |klass| klass.order.between?(first, last) }
+    end
+
+    def run
+      directory = Pathname(@options.fetch(:out, TWPipeline.work.join("run"))).tap(&:mkpath)
+      previous = nil
+
+      selection.each do |klass|
+        target = directory.join("#{klass.slug}.jsonl")
+        report(klass.slug, out: target.to_s) do
+          open_input(previous) { |input| target.open("w") { |output| build(klass).call(input, output) } }
+        end
+
+        previous = target
+      end
+    end
+
+    def open_input(path, &block)
+      path.nil? ? block.call($stdin) : path.open(&block)
+    end
+
+    def check
+      $stdin.each_line do |line|
+        text = line.strip
+        next if text.empty?
+
+        report = TWFilter.examine(TWFilter.normalize(text), policy: policy)
+        $stdout.puts(Jsonl.dump(report.to_h))
+      end
+    end
+
+    def report(label, **facts)
+      sample = Bench.measure(label, **facts) { yield }
+      warn(Jsonl.dump(sample.to_h)) unless @options[:quiet]
+      sample
+    end
+  end
+end

data/lib/twpipeline/frequency.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Frequency
+    module_function
+
+    def dispersion(counts, sizes)
+      total = counts.values.sum.to_f
+      corpus = sizes.values.sum.to_f
+      return 0.0 if total.zero? || corpus.zero?
+
+      0.5 * sizes.sum { |part, size| (counts.fetch(part, 0) / total - size / corpus).abs }
+    end
+
+    def each_key(path)
+      return to_enum(:each_key, path) unless block_given?
+
+      current = nil
+      parts = Hash.new(0)
+
+      Pathname(path).each_line do |line|
+        count, key, part = line.chomp.split("\t")
+        if key != current && !current.nil?
+          yield [current, parts]
+          parts = Hash.new(0)
+        end
+
+        current = key
+        parts[part] += count.to_i
+      end
+
+      yield [current, parts] unless current.nil?
+    end
+
+    def table(path, sizes)
+      each_key(path).map do |key, parts|
+        total = parts.values.sum
+        {
+          key: key,
+          frequency: total,
+          dispersion: dispersion(parts, sizes),
+          adjusted: (total * (1 - dispersion(parts, sizes))).round(2),
+          parts: parts.length
+        }
+      end
+    end
+  end
+end

data/lib/twpipeline/jsonl.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require "json"
+
+module TWPipeline
+  module Jsonl
+    module_function
+
+    def parse(line)
+      JSON.parse(line, symbolize_names: true)
+    rescue JSON::ParserError => error
+      raise Error, "malformed JSONL: #{error.message}"
+    end
+
+    def dump(record) = JSON.generate(record)
+
+    def each(io)
+      return to_enum(:each, io) unless block_given?
+
+      io.each_line do |line|
+        stripped = line.strip
+        yield parse(stripped) unless stripped.empty?
+      end
+    end
+
+    def wrap(text, source: "stdin", index: 0)
+      {id: "#{source}:#{index}", source: source, text: text, ok: true, findings: []}
+    end
+  end
+end

data/lib/twpipeline/parallel.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Parallel
+    module_function
+
+    def map(items, workers: TWPipeline.resources.cores, &block)
+      list = items.to_a
+      return list.map(&block) if workers < 2 || list.length < 2
+
+      slices(list, workers)
+        .map { |slice| spawn(slice, &block) }
+        .map { |pid, reader| collect(pid, reader) }
+        .flatten(1)
+    end
+
+    def fold(items, workers: TWPipeline.resources.cores, &block)
+      list = items.to_a
+      return [block.call(list)] if workers < 2 || list.length < 2
+
+      slices(list, workers)
+        .map { |slice| spawn([slice], &block) }
+        .flat_map { |pid, reader| collect(pid, reader) }
+    end
+
+    CHUNK_CAP = 1 << 30
+    CHUNK_SHARE = 32
+    LINE_CAP = 500_000
+
+    def chunk_bytes(resources = TWPipeline.resources)
+      override = TWPipeline::Resources.parse(ENV["TWP_CHUNK_BYTES"])
+      override || [resources.memory_bytes / CHUNK_SHARE, CHUNK_CAP].min
+    end
+
+    def pipe(input, output, workers: TWPipeline.resources.cores, budget: chunk_bytes, &block)
+      read = 0
+      written = 0
+      bytes = 0
+      buffer = []
+
+      flush = lambda do
+        map(buffer, workers: workers, &block).each do |row|
+          Array(row).each do |line|
+            output.puts(line)
+            written += 1
+          end
+        end
+
+        buffer = []
+        bytes = 0
+      end
+
+      lines = input.respond_to?(:each_line) ? input.each_line : input.each
+      lines.each do |line|
+        read += 1
+        buffer << line
+        bytes += line.bytesize
+        flush.call if bytes >= budget || buffer.length >= LINE_CAP
+      end
+
+      flush.call if buffer.any?
+      {read: read, written: written}
+    end
+
+    def slices(list, workers)
+      size = (list.length / workers.to_f).ceil
+      list.each_slice([size, 1].max).to_a
+    end
+
+    def spawn(slice, &block)
+      reader, writer = IO.pipe
+      pid = fork do
+        reader.close
+        writer.binmode
+        writer.write(Marshal.dump(slice.map(&block)))
+        writer.close
+        exit!(0)
+      end
+
+      writer.close
+      [pid, reader]
+    end
+
+    def collect(pid, reader)
+      reader.binmode
+      payload = reader.read
+      reader.close
+      Process.wait(pid)
+      Marshal.load(payload)
+    end
+  end
+end

data/lib/twpipeline/resources.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require "etc"
+
+module TWPipeline
+  Resources = Data.define(:cores, :memory_bytes) do
+    UNITS = {"K" => 1024, "M" => 1024 ** 2, "G" => 1024 ** 3, "T" => 1024 ** 4}.freeze
+    DEFAULT_SHARE = 0.5
+
+    def memory = format_bytes(memory_bytes)
+
+    def sort_buffer = format_bytes((memory_bytes * 0.8).to_i)
+
+    def memory_per_worker_bytes = memory_bytes / [cores, 1].max
+
+    def to_h = {cores: cores, memory_bytes: memory_bytes, memory: memory}
+
+    def to_s = "#{cores} cores, #{memory} usable"
+
+    private
+
+    def format_bytes(bytes)
+      unit, factor = UNITS.to_a.reverse.find { |_, value| bytes >= value } || ["K", 1024]
+      "#{(bytes / factor.to_f).round}#{unit}"
+    end
+
+    class << self
+      def detect(cores: nil, memory: nil)
+        new(
+          cores: (cores || ENV["TWP_CORES"] || Etc.nprocessors).to_i.clamp(1, 1024),
+          memory_bytes: parse(memory || ENV["TWP_MEMORY"]) || (total_memory * DEFAULT_SHARE).to_i
+        )
+      end
+
+      def ask(io: $stderr, input: $stdin)
+        suggested = detect
+        io.print("cores [#{suggested.cores}]: ")
+        cores = input.gets.to_s.strip
+        io.print("usable memory [#{suggested.memory}]: ")
+        memory = input.gets.to_s.strip
+
+        detect(cores: cores.empty? ? nil : cores, memory: memory.empty? ? nil : memory)
+      end
+
+      def parse(value)
+        return nil if value.nil? || value.to_s.strip.empty?
+        return value if value.is_a?(Integer)
+
+        match = value.to_s.strip.upcase.match(/\A(\d+(?:\.\d+)?)\s*([KMGT])?B?\z/)
+        raise Error, "cannot parse memory size: #{value}" if match.nil?
+
+        (match[1].to_f * UNITS.fetch(match[2], 1)).to_i
+      end
+
+      def total_memory
+        @total_memory ||= if RUBY_PLATFORM.include?("darwin")
+          `sysctl -n hw.memsize`.to_i
+        else
+          File.read("/proc/meminfo")[/MemTotal:\s+(\d+) kB/, 1].to_i * 1024
+        end
+      end
+    end
+  end
+end

data/lib/twpipeline/rusage.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require "fiddle"
+
+module TWPipeline
+  module Rusage
+    SELF = 0
+    CHILDREN = -1
+    MAXRSS_OFFSET = Fiddle::SIZEOF_LONG * 4
+    SCALE = RUBY_PLATFORM.include?("darwin") ? 1 : 1024
+
+    class << self
+      def available? = !getrusage.nil?
+
+      def max_rss(who)
+        return 0 unless available?
+
+        buffer = Fiddle::Pointer.malloc(Fiddle::SIZEOF_LONG * 32, Fiddle::RUBY_FREE)
+        return 0 unless getrusage.call(who, buffer).zero?
+
+        buffer[MAXRSS_OFFSET, Fiddle::SIZEOF_LONG].unpack1("l!") * SCALE
+      end
+
+      def peak = max_rss(SELF) + max_rss(CHILDREN)
+
+      private
+
+      def getrusage
+        return @getrusage if defined?(@getrusage)
+
+        @getrusage = Fiddle::Function.new(
+          Fiddle.dlopen(nil)["getrusage"],
+          [Fiddle::TYPE_INT, Fiddle::TYPE_VOIDP],
+          Fiddle::TYPE_INT
+        )
+      rescue Fiddle::DLError, NameError
+        @getrusage = nil
+      end
+    end
+  end
+end

data/lib/twpipeline/segmenter.rb

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+require "json"
+
+module TWPipeline
+  class Segmenter
+    MAX_WORD = 8
+    HAN_RUN = /[\u{3400}-\u{9FFF}]+/
+
+    Model = Data.define(:cost, :char_cost, :unseen_word, :unseen_char, :max_word)
+
+    class << self
+      def load(path)
+        data = JSON.parse(Pathname(path).read)
+        new(words: entries(data.fetch("words")), chars: entries(data.fetch("chars", [])))
+      end
+
+      def entries(value) = (value.is_a?(Hash) ? value.keys : Array(value)).to_set
+
+      def runs(text) = text.scan(HAN_RUN)
+    end
+
+    attr_reader :words, :model
+
+    def initialize(words:, chars: Set.new, model: nil)
+      @words = words
+      @chars = chars
+      @limit = [words.map(&:length).max.to_i, MAX_WORD].min
+      @model = model
+    end
+
+    def fit(runs, rounds: 2)
+      counts = naive_counts(runs)
+      rounds.times do
+        @model = build_model(*counts)
+        counts = recount(runs)
+      end
+
+      @model = build_model(*counts)
+      self
+    end
+
+    def call(text) = self.class.runs(text).flat_map { |run| segment(run) }
+
+    def segment(run)
+      chars = run.chars
+      size = chars.length
+      best = Array.new(size + 1, Float::INFINITY)
+      best[0] = 0.0
+      back = Array.new(size + 1, 0)
+
+      (1..size).each do |stop|
+        ([1, stop - @model.max_word + 1].max..stop).each do |start|
+          previous = best[start - 1]
+          next if previous.infinite?
+
+          token = chars[(start - 1)...stop].join
+          price = price_of(token)
+          next if price.nil?
+
+          value = previous + price
+          next unless value < best[stop]
+
+          best[stop] = value
+          back[stop] = start - 1
+        end
+      end
+
+      unwind(chars, back, size)
+    end
+
+    private
+
+    def price_of(token)
+      return @model.char_cost.fetch(token, @model.unseen_char) if token.length == 1
+      return @model.cost.fetch(token, @model.unseen_word) if @words.include?(token)
+
+      nil
+    end
+
+    def unwind(chars, back, size)
+      pieces = []
+      stop = size
+      while stop.positive?
+        start = back[stop]
+        pieces << chars[start...stop].join
+        stop = start
+      end
+
+      pieces.reverse
+    end
+
+    def naive_counts(runs)
+      word_counts = Hash.new(0)
+      char_counts = Hash.new(0)
+
+      runs.each do |run|
+        chars = run.chars
+        chars.each { |char| char_counts[char] += 1 }
+        chars.length.times do |start|
+          (2..[@limit, chars.length - start].min).each do |length|
+            token = chars[start, length].join
+            word_counts[token] += 1 if @words.include?(token)
+          end
+        end
+      end
+
+      [word_counts, char_counts]
+    end
+
+    def recount(runs)
+      word_counts = Hash.new(0)
+      char_counts = Hash.new(0)
+
+      runs.each do |run|
+        segment(run).each { |token| (token.length == 1 ? char_counts : word_counts)[token] += 1 }
+      end
+
+      [word_counts, char_counts]
+    end
+
+    def build_model(word_counts, char_counts)
+      total = [word_counts.each_value.sum + char_counts.each_value.sum, 1].max
+
+      Model.new(
+        cost: word_counts.transform_values { |count| -Math.log(count.fdiv(total)) },
+        char_cost: char_counts.transform_values { |count| -Math.log(count.fdiv(total)) },
+        unseen_word: -Math.log(0.5 / total),
+        unseen_char: -Math.log(0.2 / total),
+        max_word: @limit
+      )
+    end
+  end
+end

data/lib/twpipeline/sorting.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require "shellwords"
+
+module TWPipeline
+  module Sorting
+    module_function
+
+    def program = @program ||= ENV["TWP_SORT"] || (system("command -v gsort > /dev/null 2>&1") ? "gsort" : "sort")
+
+    def sort(*extra, resources: TWPipeline.resources)
+      tmp = TWPipeline.work.join("tmp").tap(&:mkpath)
+      [program, "-S", resources.sort_buffer, "--parallel", resources.cores.to_s, "-T", tmp.to_s, *extra].shelljoin
+    end
+
+    def unique_pipeline(floor:, resources:, collapse: false)
+      head = if collapse
+        ["LC_ALL=C #{sort("-u", resources: resources)}", "LC_ALL=C cut -f1"]
+      else
+        ["LC_ALL=C #{sort(resources: resources)}"]
+      end
+
+      [
+        *head,
+        "LC_ALL=C uniq -c",
+        "LC_ALL=C sed -E 's/^ *([0-9]+) /\\1\t/'",
+        "LC_ALL=C awk -F'\t' -v floor=#{floor.to_i} '$1 >= floor'",
+        "LC_ALL=C #{sort("-t", "\t", "-k1,1nr", "-k2,2", resources: resources)}"
+      ].join(" | ")
+    end
+
+    def tally(output_path, floor: 1, resources: TWPipeline.resources, collapse: false)
+      Pathname(output_path).dirname.mkpath
+      command = "#{unique_pipeline(floor: floor, resources: resources, collapse: collapse)} > #{output_path.to_s.shellescape}"
+      written = 0
+
+      IO.popen(command, "w") do |sink|
+        yield(
+          -> (key) {
+            sink.puts(key)
+            written += 1
+          }
+        )
+      end
+
+      raise Error, "sort pipeline failed: #{command}" unless $CHILD_STATUS.success?
+
+      {keys_written: written, rows: Pathname(output_path).each_line.count}
+    end
+  end
+end