twpipeline 0.1.0

data/lib/twpipeline/sources/hplt.rb

@@ -0,0 +1,112 @@
+# frozen_string_literal: true
+
+require "uri"
+
+module TWPipeline
+  module Sources
+    class HPLT
+      EXCLUDED_TLDS = %w[cn hk mo sg my].to_set.freeze
+      TAIWAN_TLD = "tw"
+      TAIWAN_LANG = /\Azh-(?:TW|Hant-TW)\b/i
+      URL_HEAD = 4096
+      URL = /"u":"([^"]*)"/n
+      HTML_LANG = /"html_lang":\[([^\]]*)\]/n
+      REGISTERS = %w[MT LY SP ID NA HI IN OP IP].freeze
+
+      Decision = Data.define(:keep, :gate, :host)
+
+      def initialize(allowlist: Set.new)
+        @allowlist = allowlist
+      end
+
+      def each_record(io)
+        return to_enum(:each_record, io) unless block_given?
+
+        io.each_line do |line|
+          decision = decide(line)
+          next if decision.gate == :excluded
+
+          yield build(line, decision) if decision.keep || decision.gate == :residue
+        end
+      end
+
+      def decide(line)
+        host = host_of(line)
+        return Decision.new(keep: false, gate: :excluded, host: host) if host.nil?
+
+        labels = host.split(".")
+        tld = labels.last
+        return Decision.new(keep: false, gate: :excluded, host: host) if EXCLUDED_TLDS.include?(tld)
+        return Decision.new(keep: true, gate: :tld, host: host) if tld == TAIWAN_TLD
+        return Decision.new(keep: true, gate: :allowlist, host: host) if @allowlist.include?(host)
+        return Decision.new(keep: true, gate: :html_lang, host: host) if taiwan_lang?(line)
+
+        Decision.new(keep: false, gate: :residue, host: host)
+      end
+
+      def host_of(line)
+        url = probe(line, URL_HEAD)[URL, 1] || line[URL, 1]
+        return nil if url.nil?
+
+        URI.parse(url).host&.downcase&.delete_prefix("www.")
+      rescue URI::InvalidURIError
+        url[%r{\Ahttps?://([^/:?#]+)}i, 1]&.downcase&.delete_prefix("www.")
+      end
+
+      def probe(line, size) = line.byteslice(0, size).force_encoding(Encoding::BINARY)
+
+      def taiwan_lang?(line)
+        captured = line[HTML_LANG, 1]
+        return false if captured.nil?
+
+        captured.scan(/"([^"]*)"/).flatten.any? { |tag| tag.match?(TAIWAN_LANG) }
+      end
+
+      def build(line, decision)
+        document = Jsonl.parse(line.dup.force_encoding(Encoding::UTF_8).scrub)
+        {
+          id: document[:id],
+          source: "hplt",
+          url: document[:u],
+          host: decision.host,
+          gate: decision.gate.to_s,
+          keep: decision.keep,
+          crawl: document[:crawl_id],
+          html_lang: document[:html_lang],
+          lang: Array(document[:lang]).first,
+          lang_prob: Array(document[:prob]).first,
+          doc_score: score(document),
+          register: register(document),
+          text: document[:text].to_s,
+          ok: true,
+          findings: []
+        }
+      end
+
+      def score(document)
+        scores = Array(document[:doc_scores]).grep(Numeric)
+        scores.empty? ? nil : (scores.sum / scores.length.to_f).round(2)
+      end
+
+      def register(document)
+        labels = document[:"web-register"]
+        return nil unless labels.is_a?(Hash)
+
+        labels.max_by { |_, value| value.to_f }&.first&.to_s
+      end
+
+      class << self
+        def open(path, &block)
+          command = path.to_s.end_with?(".zst") ? ["zstd", "-dc", "-T0", path.to_s] : ["cat", path.to_s]
+          IO.popen(command, "rb") { |io| block.call(io) }
+        end
+
+        def allowlist(path)
+          return Set.new if path.nil? || !Pathname(path).exist?
+
+          Pathname(path).each_line.map { |line| line.strip.downcase }.reject(&:empty?).to_set
+        end
+      end
+    end
+  end
+end

data/lib/twpipeline/stage.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  class Stage
+    REGISTRY = {}
+
+    class << self
+      def register(slug, order)
+        @slug = slug
+        @order = order
+        REGISTRY[slug] = self
+      end
+
+      def slug = @slug
+
+      def order = @order
+
+      def ordered = REGISTRY.values.sort_by(&:order)
+
+      def lookup(name)
+        REGISTRY[name] ||
+          REGISTRY.values.find { |stage| stage.slug.split("_", 2).last == name } ||
+          raise(Error, "unknown stage: #{name} (known: #{REGISTRY.keys.join(", ")})")
+      end
+    end
+
+    def initialize(policy: TWFilter::Policy.corpus, drop: false, resources: TWPipeline.resources, **options)
+      @policy = policy
+      @drop = drop
+      @resources = resources
+      @options = options
+    end
+
+    attr_reader :policy, :resources, :options
+
+    def drop? = @drop
+
+    def call(input, output) = raise NotImplementedError
+
+    private
+
+    def reject(record, findings)
+      record.merge(ok: false, findings: record.fetch(:findings, []) + findings.map(&:to_h))
+    end
+
+    def flag(record, findings)
+      return record if findings.empty?
+
+      merged = record.fetch(:findings, []) + findings.map(&:to_h)
+      record.merge(ok: record.fetch(:ok, true) && findings.none?(&:reject?), findings: merged)
+    end
+  end
+
+  class MapStage < Stage
+    def call(input, output)
+      Parallel.pipe(input, output, workers: resources.cores) { |line| process(line) }
+    end
+
+    def process(line)
+      stripped = line.strip
+      return nil if stripped.empty?
+
+      record = transform(Jsonl.parse(stripped))
+      return nil if record.nil? || (drop? && record[:ok] == false)
+
+      Jsonl.dump(record)
+    end
+
+    def transform(record) = raise NotImplementedError
+  end
+end

data/lib/twpipeline/stages/count.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Stages
+    class Count < Stage
+      register("09_count", 9)
+
+      DEFAULT_ORDERS = [1, 2, 3, 4].freeze
+
+      def call(input, output)
+        sinks = {}
+        pipelines = orders.to_h { |order| [order, sink(order, sinks)] }
+        records = 0
+
+        Jsonl.each(input) do |record|
+          records += 1
+          units = units_of(record)
+          part = part_of(record)
+          orders.each do |order|
+            emit(units, order) { |gram| pipelines.fetch(order).call(part ? "#{gram}\t#{part}" : gram) }
+          end
+        end
+
+        results = close(sinks)
+        output.puts(Jsonl.dump({records: records, orders: results}))
+        {records: records, orders: results}
+      end
+
+      private
+
+      def orders = @orders ||= Array(options.fetch(:orders, DEFAULT_ORDERS)).map(&:to_i).sort
+
+      def floor = options.fetch(:floor, 1).to_i
+
+      def unit = options.fetch(:unit, "token")
+
+      def once_per = options[:once_per]&.to_s
+
+      def part_of(record)
+        return options[:parts]&.then { |field| record[field.to_sym] } if once_per.nil?
+
+        record[once_per.to_sym] || record[:doc] || record[:id]
+      end
+
+      def directory = Pathname(options.fetch(:out, TWPipeline.work.join("counts"))).tap(&:mkpath)
+
+      def units_of(record)
+        case unit
+        in "token"
+          Array(record[:tokens])
+        in "char"
+          record.fetch(:text, "").chars
+        in other
+          raise Error, "unknown unit: #{other}"
+        end
+      end
+
+      def emit(units, order)
+        return if units.length < order
+
+        joiner = unit == "token" ? " " : ""
+        units.each_cons(order) { |gram| yield gram.join(joiner) }
+      end
+
+      def sink(order, sinks)
+        path = directory.join("#{unit}_#{order}gram.tsv")
+        holder = {path: path, thread: nil, writer: nil}
+        reader, writer = IO.pipe
+        holder[:writer] = writer
+        holder[:thread] = Thread.new do
+          Sorting.tally(path, floor: floor, resources: resources, collapse: !once_per.nil?) do |push|
+            reader.each_line { |line| push.call(line.chomp) }
+          end
+        end
+        sinks[order] = holder
+
+        -> (gram) { writer.puts(gram) }
+      end
+
+      def close(sinks)
+        sinks.transform_values do |holder|
+          holder.fetch(:writer).close
+          result = holder.fetch(:thread).value
+          {path: holder.fetch(:path).to_s}.merge(result)
+        end
+      end
+    end
+  end
+end

data/lib/twpipeline/stages/dedup.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require "digest"
+
+module TWPipeline
+  module Stages
+    class Dedup < Stage
+      register("07_dedup", 7)
+
+      MASK = (1 << 62) - 1
+
+      def call(input, output)
+        seen = Set.new
+        hosts = Hash.new { |memo, key| memo[key] = Set.new }
+        read = 0
+        written = 0
+        repeats = 0
+
+        Jsonl.each(input) do |record|
+          read += 1
+          key = fingerprint(record.fetch(:text, ""))
+          hosts[key] << record[:host] if record[:host] && hosts[key].size < ubiquity
+
+          if seen.add?(key)
+            written += 1
+            output.puts(Jsonl.dump(record))
+          else
+            repeats += 1
+            output.puts(Jsonl.dump(reject(record, [duplicate]))) unless drop?
+          end
+        end
+
+        boilerplate = hosts.count { |_, group| group.size >= ubiquity }
+        {read: read, written: written, duplicates: repeats, types: seen.size, boilerplate_types: boilerplate}
+      end
+
+      private
+
+      def ubiquity = options.fetch(:ubiquity, 5).to_i
+
+      def duplicate = TWFilter::Finding.new(check: :dedup, code: :duplicate)
+
+      def fingerprint(text) = Digest::SHA256.digest(text).unpack1("Q>") & MASK
+    end
+  end
+end

data/lib/twpipeline/stages/evidence.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Stages
+    class Evidence < Stage
+      register("06_evidence", 6)
+
+      BATCH = 20_000
+
+      def call(input, output)
+        blocks = 0
+        kept = 0
+        dropped = 0
+        batch = []
+
+        flush = lambda do
+          Parallel.map(batch, workers: resources.cores) { |records| judge(records) }.each do |verdict|
+            blocks += 1
+            verdict.each do |record|
+              record[:ok] == false ? dropped += 1 : kept += 1
+              output.puts(Jsonl.dump(record)) unless drop? && record[:ok] == false
+            end
+          end
+
+          batch = []
+        end
+
+        each_block(input) do |records|
+          batch << records
+          flush.call if batch.length >= BATCH
+        end
+
+        flush.call if batch.any?
+
+        {blocks: blocks, kept: kept, dropped: dropped}
+      end
+
+      private
+
+      def key(record) = options.fetch(:group, "doc").then { |field| record[field.to_sym] }
+
+      def size = options.fetch(:block, TWFilter::Block::SIZE).to_i
+
+      def each_block(input)
+        current = nil
+        buffer = []
+
+        Jsonl.each(input) do |record|
+          identity = key(record)
+          if (identity != current && !buffer.empty?) || buffer.length >= size
+            yield buffer
+            buffer = []
+          end
+
+          current = identity
+          buffer << record
+        end
+
+        yield buffer unless buffer.empty?
+      end
+
+      def judge(records)
+        rejected = records.count { |record| record[:ok] == false }
+        evidence = TWFilter::Evidence.count(records.map { |record| record.fetch(:text, "") }.join)
+        required = policy.evidence_per_100 * records.length / 100.0
+        contaminated = rejected > policy.block_tolerance * records.length
+        finding = block_finding(contaminated, evidence, required)
+
+        records.map { |record| finding ? flag(record, [finding]) : record.merge(evidence: evidence) }
+      end
+
+      def block_finding(contaminated, evidence, required)
+        return TWFilter::Finding.new(check: :block, code: :contaminated_block) if contaminated
+        return nil if evidence >= required
+
+        TWFilter::Finding.new(check: :block, code: :no_taiwan_evidence, detail: format("%.2f", evidence))
+      end
+    end
+  end
+end

data/lib/twpipeline/stages/ingest.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Stages
+    class Ingest < Stage
+      register("01_ingest", 1)
+
+      def call(input, output)
+        kept = 0
+        residue = 0
+        gates = Hash.new(0)
+        sink = residue_path&.then { |path|
+          path.dirname.mkpath
+          path.open("w")
+        }
+
+        per_host = Hash.new(0)
+
+        reader(input) do |io|
+          source.each_record(io) do |record|
+            gates[record[:gate]] += 1
+            next if capped?(per_host, record)
+
+            if record[:keep]
+              kept += 1
+              output.puts(Jsonl.dump(record.except(:keep)))
+            elsif sink
+              residue += 1
+              sink.puts(Jsonl.dump(record.except(:keep)))
+            end
+          end
+        end
+
+        sink&.close
+        {kept: kept, residue: residue, gates: gates}
+      end
+
+      private
+
+      def host_cap = options[:host_cap]&.to_i
+
+      def capped?(per_host, record)
+        return false if host_cap.nil? || record[:host].nil?
+
+        per_host[record[:host]] += 1
+        per_host[record[:host]] > host_cap
+      end
+
+      def source
+        @source ||= case options.fetch(:source, "hplt")
+        in "hplt"
+          Sources::HPLT.new(allowlist: Sources::HPLT.allowlist(options[:allowlist]))
+        in unknown
+          raise Error, "unknown source: #{unknown}"
+        end
+      end
+
+      def residue_path = options[:residue]&.then { |path| Pathname(path) }
+
+      def reader(input, &block)
+        path = options[:input]
+        path ? Sources::HPLT.open(path, &block) : block.call(input)
+      end
+    end
+  end
+end

data/lib/twpipeline/stages/lexicon.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Stages
+    class Lexicon < MapStage
+      register("05_lexicon", 5)
+
+      CHECKS = [TWFilter::Checks::Lexicon, TWFilter::Checks::Erhua, TWFilter::Checks::Wenyan].freeze
+
+      def transform(record)
+        subject = TWFilter::Subject.new(record.fetch(:text, ""), policy: policy)
+        flag(record, CHECKS.flat_map { |check| check.call(subject) })
+      end
+    end
+  end
+end

data/lib/twpipeline/stages/normalize.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Stages
+    class Normalize < MapStage
+      register("02_normalize", 2)
+
+      def transform(record)
+        original = record.fetch(:text, "")
+        normalized = TWFilter.normalize(original, collapse_repeats: options.fetch(:collapse_repeats, true))
+        record.merge(text: normalized, normalized: normalized != original)
+      end
+    end
+  end
+end

data/lib/twpipeline/stages/script.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Stages
+    class Script < MapStage
+      register("04_script", 4)
+
+      def transform(record)
+        subject = TWFilter::Subject.new(record.fetch(:text, ""), policy: policy)
+        flag(record.merge(tier: TWFilter::Checks::Script.tier_of(subject.text)), TWFilter::Checks::Script.call(subject))
+      end
+    end
+  end
+end

data/lib/twpipeline/stages/segment.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Stages
+    class Segment < Stage
+      register("03_segment", 3)
+
+      def call(input, output)
+        counts = Parallel.pipe(input, output, workers: resources.cores) do |line|
+          stripped = line.strip
+          next nil if stripped.empty?
+
+          split(Jsonl.parse(stripped)).map { |record| Jsonl.dump(record) }
+        end
+
+        {documents: counts.fetch(:read), sentences: counts.fetch(:written)}
+      end
+
+      private
+
+      def split(record)
+        TWFilter::Sentences
+          .split(record.fetch(:text, ""), clause: options.fetch(:clause, true))
+          .each_with_index
+          .filter_map { |text, index| sentence(record, text, index) }
+      end
+
+      CARRIED = %i[source host gate register].freeze
+
+      def sentence(record, text, index)
+        subject = TWFilter::Subject.new(text, policy: policy)
+        findings = TWFilter::Checks::Shape.call(subject)
+        return nil if drop? && findings.any?(&:reject?)
+
+        flag(
+          record.slice(*CARRIED).merge(
+            id: "#{record[:id]}##{index}",
+            doc: record[:id],
+            index: index,
+            han: subject.han,
+            text: text,
+            ok: true,
+            findings: []
+          ),
+          findings
+        )
+      end
+    end
+  end
+end

data/lib/twpipeline/stages/tokenize.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  module Stages
+    class Tokenize < MapStage
+      register("08_tokenize", 8)
+
+      FIT_LINES = 200_000
+
+      def call(input, output)
+        head = []
+        lines = input.each_line
+        lines.each { |line|
+          head << line
+          break if head.length >= fit_lines
+        }
+        prime(head)
+
+        Parallel.pipe(head.each + lines, output, workers: resources.cores) { |line| process(line) }
+      end
+
+      def transform(record) = record.merge(tokens: segmenter.call(record.fetch(:text, "")))
+
+      private
+
+      def fit_lines = options.fetch(:fit_lines, FIT_LINES).to_i
+
+      def segmenter = @segmenter ||= Segmenter.load(vocabulary)
+
+      def vocabulary
+        options[:vocabulary] ||
+          ENV["TWP_VOCABULARY"] ||
+          raise(Error, "--vocabulary is required (JSON with \"words\" and \"chars\")")
+      end
+
+      def prime(head)
+        sample = head.filter_map { |line| Jsonl.parse(line.strip)[:text] unless line.strip.empty? }
+        segmenter.fit(Segmenter.runs(sample.join("\n")), rounds: options.fetch(:rounds, 2).to_i)
+      end
+    end
+  end
+end

data/lib/twpipeline/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module TWPipeline
+  VERSION = "0.1.0"
+end

data/lib/twpipeline.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require "English"
+require "pathname"
+require "set"
+require "twfilter"
+
+require_relative "twpipeline/version"
+require_relative "twpipeline/resources"
+require_relative "twpipeline/parallel"
+require_relative "twpipeline/rusage"
+require_relative "twpipeline/bench"
+require_relative "twpipeline/jsonl"
+require_relative "twpipeline/sorting"
+require_relative "twpipeline/frequency"
+require_relative "twpipeline/segmenter"
+require_relative "twpipeline/stage"
+require_relative "twpipeline/sources/hplt"
+require_relative "twpipeline/stages/ingest"
+require_relative "twpipeline/stages/normalize"
+require_relative "twpipeline/stages/segment"
+require_relative "twpipeline/stages/script"
+require_relative "twpipeline/stages/lexicon"
+require_relative "twpipeline/stages/evidence"
+require_relative "twpipeline/stages/dedup"
+require_relative "twpipeline/stages/tokenize"
+require_relative "twpipeline/stages/count"
+
+module TWPipeline
+  class Error < StandardError
+  end
+
+  ROOT = Pathname(__dir__).parent
+
+  class << self
+    attr_writer :resources
+
+    def resources = @resources ||= Resources.detect
+
+    def work = Pathname(ENV.fetch("TWP_WORK", ROOT.parent.join("work").to_s))
+
+    def data(*parts) = ROOT.join("data", *parts)
+
+    def warm_up
+      TWFilter::Checks::Script.tiers
+      TWFilter::Checks::Lexicon.hard_terms
+      TWFilter::Checks::Lexicon.soft_terms
+      TWFilter::Checks::Lexicon.exceptions
+      TWFilter::Checks::Lexicon.regional_terms
+      TWFilter::Checks::Lexicon.foreign_topics
+      TWFilter::Checks::Lexicon.cantonese_chars
+      TWFilter::Checks::Erhua.headed
+      TWFilter::Checks::Erhua.tailed
+      TWFilter::Checks::Wenyan.markers
+      TWFilter::Evidence.terms
+      self
+    end
+  end
+end