gigatoken 0.2.2 → 0.4.0

data/lib/gigatoken/cli/support.rb

@@ -4,16 +4,29 @@ require "etc"
 
 module Gigatoken
   module CLI
-    # Helpers shared by the bench and validate commands: tokenizer loading,
-    # byte-size parsing, document splitting, and CPU identification.
+    # Helpers shared by the bench and validate commands: argument checking,
+    # tokenizer loading, byte-size parsing, document splitting, and CPU
+    # identification.
     module Support
       SIZE_UNITS = {"" => 1, "k" => 10**3, "m" => 10**6, "g" => 10**9, "t" => 10**12}.freeze
       private_constant :SIZE_UNITS
 
-      SIZE_PATTERN = /\A\s*(\d+(?:\.\d+)?)\s*([kmgt]?)i?b?\s*\z/i
+      # "KiB" and friends are binary, as everywhere else that spells the i.
+      BINARY_SIZE_UNITS = {"" => 1, "k" => 2**10, "m" => 2**20, "g" => 2**30, "t" => 2**40}.freeze
+      private_constant :BINARY_SIZE_UNITS
+
+      SIZE_PATTERN = /\A\s*(\d+(?:\.\d+)?)\s*([kmgt]?)(i?)b?\s*\z/i
       private_constant :SIZE_PATTERN
 
       class << self
+        # The argument shapes dry-cli itself lets through: FILES declared
+        # `required: true` still arrives empty, and an empty separator would
+        # split every byte into its own document.
+        def check_usage!(files, separator)
+          raise Gigatoken::Error, "FILES is required: name at least one file to encode" if files.empty?
+          raise Gigatoken::Error, "--doc-separator cannot be empty" if separator == ""
+        end
+
         # Load TOKENIZER: a tokenizer.json path/directory, a packaged
         # tiktoken encoding name, a HuggingFace repo id, or a .tiktoken file
         # — see Gigatoken::Tokenizer.load. `pretokenizer:` is forwarded
@@ -23,15 +36,17 @@ module Gigatoken
           Gigatoken::Tokenizer.load(spec, pretokenizer: pretokenizer)
         end
 
-        # Parse a decimal byte size like "100MB", "2.5GB", or "1000000";
-        # "none"/"unlimited" means no limit.
+        # Parse a byte size like "100MB", "2.5GB", "64KiB" or "1000000" —
+        # decimal units unless the prefix spells the i, which makes it
+        # binary; "none"/"unlimited" means no limit.
         def parse_size(text)
           return nil if ["none", "unlimited"].include?(text.strip.downcase)
 
           match = SIZE_PATTERN.match(text)
           raise Gigatoken::Error, "cannot parse size #{text.inspect}; expected something like 100MB (or 'none')" unless match
 
-          (match[1].to_f * SIZE_UNITS.fetch(match[2].downcase)).to_i
+          units = match[3].empty? ? SIZE_UNITS : BINARY_SIZE_UNITS
+          (match[1].to_f * units.fetch(match[2].downcase)).to_i
         end
 
         # A Native::TextFileSource for FILES, split on `separator` when
@@ -42,15 +57,23 @@ module Gigatoken
 
         # Whole files as raw bytes, one document per file, or (with a
         # separator) the separator-split pieces of each file in order, empty
-        # documents skipped.
+        # documents skipped. Compressed files are decompressed first, as the
+        # native file sources do, so both paths split the same bytes.
         def split_docs(files, separator)
-          raws = files.map { |file| File.binread(file.to_s) }
+          raws = files.map { |file| read_decompressed(file.to_s) }
           return raws if separator.nil?
 
           sep = separator.b
           raws.flat_map { |raw| raw.split(sep).reject(&:empty?) }
         end
 
+        # The bytes the tokenizer actually sees for FILES — a compressed
+        # file's decompressed size, not its size on disk — so throughput is
+        # reported over the input that was tokenized.
+        def input_bytesize(files)
+          files.sum { |file| read_decompressed(file.to_s).bytesize }
+        end
+
         # The prefix of `docs` totalling at most `limit_bytes`, byte-
         # truncating the final document to fill the budget. Unlike a
         # text-comparison tool, gigatoken encodes raw bytes and does not
@@ -90,6 +113,15 @@ module Gigatoken
 
         private
 
+        # One file's bytes, decompressed. `Native.read_input` is the core's
+        # own decoder — the one the native file sources load through — so
+        # `.gz`, `.zst`/`.zstd` and plain files are detected and read here
+        # exactly as they are on the other side of `validate`, and an
+        # unreadable file raises Gigatoken::Error naming it.
+        def read_decompressed(path)
+          Gigatoken::Native.read_input(path)
+        end
+
         def darwin_cpu_info
           name = sysctl("machdep.cpu.brand_string")
           [name, sysctl_int("hw.physicalcpu"), sysctl_int("hw.packages")]

data/lib/gigatoken/cli/validate.rb

@@ -20,6 +20,7 @@ module Gigatoken
       option :pretokenizer, desc: "pretokenizer scheme, required when TOKENIZER is a .tiktoken file (one of #{Native.pretokenizer_names.join(", ")}); ignored otherwise"
 
       def call(tokenizer:, files:, doc_separator: nil, pretokenizer: nil, **)
+        Support.check_usage!(files, doc_separator)
         gt_tokenizer = Support.load_tokenizer(tokenizer, pretokenizer: pretokenizer)
 
         via_files = gt_tokenizer.encode_files(Support.text_file_source(files, doc_separator))
@@ -38,7 +39,7 @@ module Gigatoken
         end
 
         out.puts "validation OK: #{via_files.length} documents match"
-      rescue Gigatoken::Error => e
+      rescue Gigatoken::Error, SystemCallError => e
         err.puts "error: #{e.message}"
         exit(1)
       end

data/lib/gigatoken/encodings.rb

@@ -40,6 +40,10 @@ module Gigatoken
     HARMONY_RESERVED_TAIL = (200013..201087).to_h { |id| ["<|reserved_#{id}|>", id] }.freeze
     private_constant :HARMONY_RESERVED_TAIL
 
+    # Deep-frozen: entries, their `special_tokens` tables and the `rank_file`
+    # paths. `Tokenizer#special_tokens` hands the registry's own Hash back to
+    # callers, so anything less lets one caller's poke rewrite what every
+    # later `from_encoding` in the process loads.
     REGISTRY = {
       "r50k_base" => {
         rank_file: File.join(DATA_DIR, "r50k_base.tiktoken"),
@@ -67,7 +71,7 @@ module Gigatoken
         pretokenizer: "o200k",
         special_tokens: HARMONY_HEAD_TOKENS.merge(HARMONY_RESERVED_TAIL).freeze
       }
-    }.freeze
+    }.each_value { |encoding| encoding.each_value(&:freeze).freeze }.freeze
     private_constant :REGISTRY
 
     # The packaged encoding names — the single source error messages naming
@@ -88,15 +92,16 @@ module Gigatoken
 
     class << self
       # The {rank_file:, pretokenizer:, special_tokens:} registered for a
-      # packaged encoding name, or nil.
+      # packaged encoding name, or nil. Names are Strings or Symbols, as
+      # Tokenizer.load accepts both.
       def [](name)
-        REGISTRY[name]
+        REGISTRY[name.to_s]
       end
 
       # Why `name` can't be packaged, or nil when there's no reason on
       # record (it's either packaged, or simply not one gigatoken knows of).
       def unpackable_reason(name)
-        UNPACKABLE_REASONS[name]
+        UNPACKABLE_REASONS[name.to_s]
       end
     end
   end

data/lib/gigatoken/hub.rb

@@ -2,6 +2,7 @@
 
 require "async"
 require "async/http"
+require "async/http/proxy"
 require "fileutils"
 require "pathname"
 
@@ -32,6 +33,39 @@ module Gigatoken
     MAX_REDIRECTS = 10
     private_constant :MAX_REDIRECTS
 
+    # Connect/read timeout in seconds, and the bound on the request phase —
+    # huggingface_hub's HF_HUB_ETAG_TIMEOUT / HF_HUB_DOWNLOAD_TIMEOUT
+    # default.
+    DEFAULT_TIMEOUT = 10
+    private_constant :DEFAULT_TIMEOUT
+
+    # A "." or ".." path segment: the traversal that must never reach a URL
+    # or the cache, wherever it comes from.
+    DOT_SEGMENT = /\A\.\.?\z/
+    private_constant :DOT_SEGMENT
+
+    # Everything outside RFC 3986's unreserved set, which is what
+    # huggingface_hub's `quote` percent-encodes.
+    RESERVED = /[^A-Za-z0-9\-._~]/
+    private_constant :RESERVED
+
+    # How a request, or the body read that follows it, fails underneath
+    # async-http: a connect/read timeout, a refused or reset connection, a
+    # proxy refusing the CONNECT tunnel, DNS resolution, a body cut short, a
+    # malformed endpoint, TLS. All of them reach the caller as
+    # Gigatoken::HubError.
+    TRANSPORT_ERRORS = [
+      Async::TimeoutError,
+      Async::HTTP::Proxy::ConnectFailure,
+      IOError,
+      SocketError,
+      SystemCallError,
+      URI::InvalidURIError,
+      Protocol::HTTP::Error,
+      OpenSSL::SSL::SSLError
+    ].freeze
+    private_constant :TRANSPORT_ERRORS
+
     class << self
       # Whether `name` is shaped like a HuggingFace Hub repo id: `org/name`,
       # or a bare legacy repo name like `gpt2`. At most one slash, and not
@@ -99,15 +133,50 @@ module Gigatoken
         revision.match?(/\A[0-9a-f]{40}\z/)
       end
 
+      # The Hub endpoint, resolved like huggingface_hub does it: HF_ENDPOINT,
+      # then https://huggingface.co.
+      def default_endpoint
+        env("HF_ENDPOINT") || DEFAULT_ENDPOINT
+      end
+
+      # Whether `value` is usable as a URL path and a cache path component:
+      # not empty, not absolute, no NUL byte, no "." or ".." segment.
+      # huggingface_hub rejects the same traversals in validate_repo_id.
+      def safe_component?(value)
+        !value.empty? && !value.include?("\0") && !value.start_with?("/") &&
+          value.split("/", -1).none? { |segment| segment.match?(DOT_SEGMENT) }
+      end
+
+      # The proxy URL the environment names for a `scheme` request to
+      # `hostname`, or nil: http_proxy/https_proxy, lowercase spelling
+      # first, suppressed by a matching no_proxy entry — what requests does,
+      # and so what huggingface_hub inherits.
+      def proxy_url(scheme, hostname)
+        url = env("#{scheme}_proxy") || env("#{scheme.upcase}_PROXY")
+        url unless url.nil? || no_proxy?(hostname)
+      end
+
       private
 
+      # no_proxy is a comma-separated list of host suffixes, or "*" for
+      # everything.
+      def no_proxy?(hostname)
+        host = hostname.downcase
+        (env("no_proxy") || env("NO_PROXY")).to_s.split(",").any? do |entry|
+          entry = entry.strip.downcase.delete_prefix(".")
+          next true if entry == "*"
+
+          !entry.empty? && (host == entry || host.end_with?(".#{entry}"))
+        end
+      end
+
       def env(key)
         value = ENV[key]
         value unless value.nil? || value.empty?
       end
 
       def word_part?(part, first_alnum:)
-        return false if part.nil? || part.empty?
+        return false if part.nil? || part.empty? || part.match?(DOT_SEGMENT)
 
         first_ok = first_alnum ? part[0].match?(/[A-Za-z0-9]/) : word_char?(part[0])
         first_ok && part[1..].chars.all? { |c| word_char?(c) }
@@ -125,15 +194,28 @@ module Gigatoken
 
     # @parameter endpoint [String] the Hub endpoint to fetch from — override
     #   for pointing at a local server in tests (dependency injection, not a
-    #   mock).
-    def initialize(endpoint: DEFAULT_ENDPOINT)
+    #   mock); defaults to Hub.default_endpoint (HF_ENDPOINT, then
+    #   huggingface.co).
+    # @parameter timeout [Numeric] connect/read timeout in seconds, applied
+    #   to every request and to the request phase as a whole.
+    def initialize(endpoint: self.class.default_endpoint, timeout: DEFAULT_TIMEOUT)
       @endpoint = endpoint.chomp("/")
-      @internet = Async::HTTP::Internet.new
+      @timeout = timeout
     end
 
     # Path of `filename` from Hub repo `repo_id` at `revision`, served from
-    # the standard HF cache, downloading into it first when absent.
+    # the standard HF cache, downloading into it first when absent. The
+    # three caller-supplied components are checked before any request or
+    # filesystem access: one of them carrying `..` would otherwise read and
+    # overwrite files outside the cache.
     def hub_file(repo_id, filename = "tokenizer.json", revision: "main")
+      {"repo id" => repo_id, "filename" => filename, "revision" => revision}.each do |what, value|
+        next if self.class.safe_component?(value)
+
+        raise HubError, "#{what} #{value.inspect}: must not be empty or absolute, " \
+          "or contain a NUL byte or a \".\" or \"..\" path segment"
+      end
+
       self.class.cached_file(repo_id, filename, revision) ||
         Sync { fetch(repo_id, filename, revision) }
     end
@@ -145,31 +227,127 @@ module Gigatoken
     # recording the branch ref so later lookups (ours and
     # huggingface_hub's) resolve it.
     def fetch(repo_id, filename, revision)
-      url = "#{@endpoint}/#{repo_id}/resolve/#{revision}/#{filename}"
+      clients = []
+      url = resolve_url(repo_id, filename, revision)
       token = self.class.hf_token
-      response = @internet.get(url, auth_headers(token))
+      headers = auth_headers(token)
+      response = get(url, headers, clients)
       # Unlisted headers parse as a Header::Generic (an Array of values);
       # x-repo-commit is always a single value, so flatten it to a String.
       commit = response.headers["x-repo-commit"]&.to_s
 
-      # Redirects are followed by hand: resolve/ URLs answer with the
-      # x-repo-commit header and a redirect to a CDN for LFS files, and the
-      # Authorization header must not travel to the other host.
+      # Redirects are followed by hand: a resolve/ URL answers an LFS file
+      # with a redirect to a CDN, a renamed repo with one to its new name,
+      # and both headers below need a rule of their own across the hop.
       hops = 0
       while (300...400).cover?(response.status)
-        ensure_ok!(url, response.status, !!token)
-        hops += 1
-        raise Error, "#{url}: too many redirects" if hops > MAX_REDIRECTS
-
-        location = response.headers["location"] ||
-          raise(Error, "#{url}: redirect with no Location header")
+        location = response.headers["location"]
         response.close
-        url = absolutize(location, url)
-        response = @internet.get(url, {"user-agent" => "gigatoken"})
+        raise HubError, "#{url}: redirect with no Location header" unless location
+        raise HubError, "#{url}: too many redirects" if (hops += 1) > MAX_REDIRECTS
+
+        target = absolutize(location, url)
+        # A renamed repo answers with a same-origin redirect that still needs
+        # the token; the LFS CDN elsewhere authenticates by signed URL and
+        # must never see it — huggingface_hub drops the header on the same
+        # rule.
+        headers = auth_headers(nil) unless same_origin?(target, url)
+        url = target
+        response = get(url, headers, clients)
+        # An LFS file carries x-repo-commit on the resolve/ hop; a renamed
+        # repo carries it only on the hop that finally answers 200.
+        commit ||= response.headers["x-repo-commit"]&.to_s
+      end
+      ensure_ok!(url, response, !!token)
+      ensure_commit!(url, commit)
+
+      write_to_cache(repo_id, filename, revision, commit, response)
+    rescue *TRANSPORT_ERRORS => e
+      raise HubError, "#{url}: #{e.message} (#{e.class})"
+    ensure
+      close_all(response, clients)
+    end
+
+    # Release everything the fetch holds, in the one place every exit passes
+    # through. Order matters: an unread body keeps its connection checked
+    # out, and Async::HTTP::Client#close waits for its pool to drain, so the
+    # response goes first and the clients innermost-first — a tunnel client's
+    # connection is held by the proxy client opened under it. Bounded, and
+    # deaf to the ways a broken connection fails to close: a connection some
+    # failure left mid-stream would otherwise stall the drain, and with it
+    # the Sync this all runs inside, for good. A leaked socket is worth less
+    # than the caller's result or exception.
+    def close_all(response, clients)
+      Async::Task.current.with_timeout(@timeout) do
+        response&.close
+        clients.reverse_each(&:close)
       end
-      ensure_ok!(url, response.status, !!token)
+    rescue *TRANSPORT_ERRORS
+      nil
+    end
+
+    # `endpoint/repo/resolve/revision/filename`, percent-encoded the way
+    # huggingface_hub's `quote` does it: `revision` whole (`safe=""`), so
+    # `refs/pr/1` travels as `refs%2Fpr%2F1`, while the repo id and the
+    # filename keep their slashes.
+    def resolve_url(repo_id, filename, revision)
+      "#{@endpoint}/#{escape_path(repo_id)}/resolve/#{escape(revision)}/#{escape_path(filename)}"
+    end
 
-      write_to_cache(repo_id, filename, revision, commit || revision, response)
+    def escape_path(value)
+      value.split("/", -1).map { |segment| escape(segment) }.join("/")
+    end
+
+    def escape(value)
+      value.gsub(RESERVED) { |char| char.bytes.map { |byte| format("%%%02X", byte) }.join }
+    end
+
+    # GET `url`, through the proxy the environment names for it when there
+    # is one: an http request travels to the proxy with the absolute URI in
+    # its request line, an https one through a CONNECT tunnel — how requests,
+    # and so huggingface_hub, routes them. The clients stay open (the body
+    # is still to be streamed); #fetch closes them.
+    def get(url, headers, clients)
+      endpoint = endpoint_for(url)
+      proxy = self.class.proxy_url(endpoint.scheme, endpoint.hostname)
+      proxy &&= endpoint_for(proxy)
+
+      client =
+        if proxy.nil?
+          open_client(endpoint, clients)
+        elsif endpoint.secure?
+          # https tunnels through the proxy with CONNECT, then speaks TLS to
+          # the origin as if it had connected to it directly.
+          open_client(open_client(proxy, clients).proxied_endpoint(endpoint), clients)
+        else
+          open_client(proxy, clients)
+        end
+      # An http proxy is addressed with the absolute URI in the request line;
+      # a direct or tunnelled request carries just the path.
+      target = (proxy && !endpoint.secure?) ? url : endpoint.path
+
+      request = Protocol::HTTP::Request["GET", target, headers, scheme: endpoint.scheme, authority: endpoint.authority]
+      # A CONNECT tunnel's socket carries no timeout of its own, so the
+      # request phase is bounded here whichever way it was routed; the body
+      # then streams under the endpoint's own connect/read timeout.
+      Async::Task.current.with_timeout(@timeout) { client.call(request) }
+    end
+
+    # A malformed URL comes back as URI::InvalidURIError, one that cannot be
+    # routed (no scheme or host) as ArgumentError; both mean the same thing
+    # to the caller, and neither is worth a backtrace.
+    def endpoint_for(url)
+      Async::HTTP::Endpoint.parse(url, timeout: @timeout)
+    rescue ArgumentError => e
+      raise HubError, "#{url}: #{e.message} (#{e.class})"
+    end
+
+    # A client for one hop, remembered in `clients` so #fetch can close it
+    # once the body is written. retries: 1 — a client built for this hop has
+    # no stale pooled connection for a retry to rescue, and async-http's
+    # default of 3 would pay @timeout three times over.
+    def open_client(endpoint, clients)
+      Async::HTTP::Client.new(endpoint, retries: 1).tap { |client| clients << client }
     end
 
     def auth_headers(token)
@@ -192,51 +370,75 @@ module Gigatoken
       rescue
         FileUtils.rm_f(tmp)
         raise
-      ensure
-        response.close
       end
       File.rename(tmp, target)
 
-      if !self.class.commit_hash?(revision) && revision != commit
-        refs_dir = repo_dir.join("refs")
-        FileUtils.mkdir_p(refs_dir)
-        File.write(refs_dir.join(revision), commit)
+      # A nested revision like refs/pr/1 is a nested ref file, so its parent
+      # has to exist — huggingface_hub mkdir -p's the same path.
+      if revision != commit
+        ref_path = repo_dir.join("refs", revision)
+        FileUtils.mkdir_p(ref_path.dirname)
+        File.write(ref_path, commit)
       end
 
       target
     end
 
-    def ensure_ok!(url, status, had_token)
+    # Raise on a non-success status. The unread body is left to #close_all,
+    # which every exit from #fetch passes through: it keeps the HTTP/1.x
+    # connection — and the Sync around it — alive until it is closed, so the
+    # exception would otherwise never reach the caller.
+    def ensure_ok!(url, response, had_token)
+      status = response.status
       return if (200...400).cover?(status)
 
       case status
       when 404
-        raise Error, "#{url}: HTTP 404 — no such repo with that file, and no such local file either"
+        raise HubError, "#{url}: HTTP 404 — no such repo with that file, and no such local file either"
       when 401, 403
         token_note = had_token ? "the request used the discovered token" : "no token was found"
-        raise Error,
+        raise HubError,
           "#{url}: HTTP #{status} — the repo may be private or gated (#{token_note}; set HF_TOKEN or run " \
           "`hf auth login`, and accept the repo's terms on huggingface.co if it is gated)"
       else
-        raise Error, "#{url}: HTTP #{status}"
+        raise HubError, "#{url}: HTTP #{status}"
       end
     end
 
+    # The snapshot directory is named by a server-controlled header, so only
+    # a real commit hash may become one: anything else would write the body
+    # wherever the header points. A missing header also means the endpoint
+    # isn't a Hub, whose snapshot would never be found again.
+    def ensure_commit!(url, commit)
+      return if self.class.commit_hash?(commit.to_s)
+
+      raise HubError, "#{url}: response is missing a usable x-repo-commit header (#{commit.inspect}) — it does not " \
+        "seem to be served by a HuggingFace Hub endpoint; if HF_ENDPOINT is set, check that it points to a " \
+        "Hub-compatible endpoint, and otherwise check your firewall and proxy settings"
+    end
+
     # A redirect Location resolved against the request URL: absolute URLs
     # pass through, host-relative (`/x/y`) and path-relative ones join the
     # base.
     def absolutize(location, base)
       return location if location.include?("://")
 
-      origin_end = base.index("://") ? base.index("://") + 3 : 0
-      origin_end = base.index("/", origin_end) || base.length
+      base_origin = origin(base)
+      return "#{base_origin}#{location}" if location.start_with?("/")
 
-      if location.start_with?("/")
-        "#{base[0...origin_end]}#{location}"
-      else
-        dir_end = base.rindex("/") || base.length
-        "#{base[0...[dir_end, origin_end].max]}/#{location}"
-      end
+      dir_end = base.rindex("/") || base.length
+      "#{base[0...[dir_end, base_origin.length].max]}/#{location}"
+    end
+
+    # Scheme and authority of a URL — everything before its path. Two URLs
+    # sharing one may pass the Authorization header between them.
+    def origin(url)
+      scheme_end = url.index("://")
+      url[0...(url.index("/", scheme_end ? scheme_end + 3 : 0) || url.length)]
+    end
+
+    def same_origin?(url, other)
+      origin(url).casecmp?(origin(other))
     end
   end
 end

data/lib/gigatoken/packed_result.rb

@@ -28,9 +28,17 @@ module Gigatoken
       lens.sum
     end
 
-    # Array of token ids for document `i`, materialized on demand.
-    def [](i)
-      buffer.get_values(Array.new(lens[i], :u32), @offsets[i] * 4)
+    # Array of token ids for document `i`, materialized on demand; negative
+    # indices count from the end, out-of-range ones give nil, and a
+    # non-Integer index is a TypeError, all like Array.
+    def [](index)
+      i = Integer.try_convert(index)
+      raise TypeError, "no implicit conversion of #{index.class} into Integer" if i.nil?
+
+      i += size if i.negative?
+      return unless i >= 0 && i < size
+
+      buffer.values(:u32, @offsets[i] * 4, lens[i])
     end
 
     def each
@@ -42,7 +50,7 @@ module Gigatoken
     # A ragged Array of Arrays, one per document — the same shape
     # `encode_batch`/`encode_files` return with `packed: false`.
     def to_a
-      each.to_a
+      Array.new(size) { |i| self[i] }
     end
   end
 end