gigatoken 0.3.0 → 0.4.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -20,30 +20,36 @@ module Gigatoken
20
20
  option :pretokenizer, desc: "pretokenizer scheme, required when TOKENIZER is a .tiktoken file (one of #{Native.pretokenizer_names.join(", ")}); ignored otherwise"
21
21
 
22
22
  def call(tokenizer:, files:, doc_separator: nil, limit_bytes: "none", parallel: true, packed: false, pretokenizer: nil, **)
23
+ Support.check_usage!(files, doc_separator)
23
24
  limit = Support.parse_size(limit_bytes)
24
25
  out.puts "#{label("cpu")}: #{Support.cpu_info}"
25
26
 
26
27
  gt_tokenizer = Support.load_tokenizer(tokenizer, pretokenizer: pretokenizer)
27
28
 
29
+ # Only the batch path materializes the documents; the native paths
30
+ # leave this nil and count their bytes off the clock below.
31
+ docs = nil
28
32
  start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
29
33
  if packed
30
34
  encoded = gt_tokenizer.encode_files(Support.text_file_source(files, doc_separator), parallel: parallel, packed: true)
31
- n_bytes = files.sum { |file| File.size(file) }
32
35
  n_tokens = encoded.token_count
33
36
  elsif parallel
34
37
  docs = Support.subset_docs(Support.split_docs(files, doc_separator), limit)
35
38
  encoded = gt_tokenizer.encode_batch(docs)
36
- n_bytes = docs.sum(&:bytesize)
37
39
  n_tokens = encoded.sum(&:length)
38
40
  else
39
41
  encoded = gt_tokenizer.encode_files(Support.text_file_source(files, doc_separator), parallel: false)
40
- n_bytes = files.sum { |file| File.size(file) }
41
42
  n_tokens = encoded.sum(&:length)
42
43
  end
43
44
  seconds = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start
44
45
 
46
+ # The throughput is over the bytes encoded, which for a compressed
47
+ # file are its decompressed bytes. Counting those means reading the
48
+ # file, so the native paths count off the clock; the batch path
49
+ # already holds the documents it encoded.
50
+ n_bytes = docs ? docs.sum(&:bytesize) : Support.input_bytesize(files)
45
51
  out.puts report("gigatoken", seconds, n_bytes, n_tokens)
46
- rescue Gigatoken::Error => e
52
+ rescue Gigatoken::Error, SystemCallError => e
47
53
  err.puts "error: #{e.message}"
48
54
  exit(1)
49
55
  end
@@ -4,16 +4,29 @@ require "etc"
4
4
 
5
5
  module Gigatoken
6
6
  module CLI
7
- # Helpers shared by the bench and validate commands: tokenizer loading,
8
- # byte-size parsing, document splitting, and CPU identification.
7
+ # Helpers shared by the bench and validate commands: argument checking,
8
+ # tokenizer loading, byte-size parsing, document splitting, and CPU
9
+ # identification.
9
10
  module Support
10
11
  SIZE_UNITS = {"" => 1, "k" => 10**3, "m" => 10**6, "g" => 10**9, "t" => 10**12}.freeze
11
12
  private_constant :SIZE_UNITS
12
13
 
13
- SIZE_PATTERN = /\A\s*(\d+(?:\.\d+)?)\s*([kmgt]?)i?b?\s*\z/i
14
+ # "KiB" and friends are binary, as everywhere else that spells the i.
15
+ BINARY_SIZE_UNITS = {"" => 1, "k" => 2**10, "m" => 2**20, "g" => 2**30, "t" => 2**40}.freeze
16
+ private_constant :BINARY_SIZE_UNITS
17
+
18
+ SIZE_PATTERN = /\A\s*(\d+(?:\.\d+)?)\s*([kmgt]?)(i?)b?\s*\z/i
14
19
  private_constant :SIZE_PATTERN
15
20
 
16
21
  class << self
22
+ # The argument shapes dry-cli itself lets through: FILES declared
23
+ # `required: true` still arrives empty, and an empty separator would
24
+ # split every byte into its own document.
25
+ def check_usage!(files, separator)
26
+ raise Gigatoken::Error, "FILES is required: name at least one file to encode" if files.empty?
27
+ raise Gigatoken::Error, "--doc-separator cannot be empty" if separator == ""
28
+ end
29
+
17
30
  # Load TOKENIZER: a tokenizer.json path/directory, a packaged
18
31
  # tiktoken encoding name, a HuggingFace repo id, or a .tiktoken file
19
32
  # — see Gigatoken::Tokenizer.load. `pretokenizer:` is forwarded
@@ -23,15 +36,17 @@ module Gigatoken
23
36
  Gigatoken::Tokenizer.load(spec, pretokenizer: pretokenizer)
24
37
  end
25
38
 
26
- # Parse a decimal byte size like "100MB", "2.5GB", or "1000000";
27
- # "none"/"unlimited" means no limit.
39
+ # Parse a byte size like "100MB", "2.5GB", "64KiB" or "1000000" —
40
+ # decimal units unless the prefix spells the i, which makes it
41
+ # binary; "none"/"unlimited" means no limit.
28
42
  def parse_size(text)
29
43
  return nil if ["none", "unlimited"].include?(text.strip.downcase)
30
44
 
31
45
  match = SIZE_PATTERN.match(text)
32
46
  raise Gigatoken::Error, "cannot parse size #{text.inspect}; expected something like 100MB (or 'none')" unless match
33
47
 
34
- (match[1].to_f * SIZE_UNITS.fetch(match[2].downcase)).to_i
48
+ units = match[3].empty? ? SIZE_UNITS : BINARY_SIZE_UNITS
49
+ (match[1].to_f * units.fetch(match[2].downcase)).to_i
35
50
  end
36
51
 
37
52
  # A Native::TextFileSource for FILES, split on `separator` when
@@ -42,15 +57,23 @@ module Gigatoken
42
57
 
43
58
  # Whole files as raw bytes, one document per file, or (with a
44
59
  # separator) the separator-split pieces of each file in order, empty
45
- # documents skipped.
60
+ # documents skipped. Compressed files are decompressed first, as the
61
+ # native file sources do, so both paths split the same bytes.
46
62
  def split_docs(files, separator)
47
- raws = files.map { |file| File.binread(file.to_s) }
63
+ raws = files.map { |file| read_decompressed(file.to_s) }
48
64
  return raws if separator.nil?
49
65
 
50
66
  sep = separator.b
51
67
  raws.flat_map { |raw| raw.split(sep).reject(&:empty?) }
52
68
  end
53
69
 
70
+ # The bytes the tokenizer actually sees for FILES — a compressed
71
+ # file's decompressed size, not its size on disk — so throughput is
72
+ # reported over the input that was tokenized.
73
+ def input_bytesize(files)
74
+ files.sum { |file| read_decompressed(file.to_s).bytesize }
75
+ end
76
+
54
77
  # The prefix of `docs` totalling at most `limit_bytes`, byte-
55
78
  # truncating the final document to fill the budget. Unlike a
56
79
  # text-comparison tool, gigatoken encodes raw bytes and does not
@@ -90,6 +113,15 @@ module Gigatoken
90
113
 
91
114
  private
92
115
 
116
+ # One file's bytes, decompressed. `Native.read_input` is the core's
117
+ # own decoder — the one the native file sources load through — so
118
+ # `.gz`, `.zst`/`.zstd` and plain files are detected and read here
119
+ # exactly as they are on the other side of `validate`, and an
120
+ # unreadable file raises Gigatoken::Error naming it.
121
+ def read_decompressed(path)
122
+ Gigatoken::Native.read_input(path)
123
+ end
124
+
93
125
  def darwin_cpu_info
94
126
  name = sysctl("machdep.cpu.brand_string")
95
127
  [name, sysctl_int("hw.physicalcpu"), sysctl_int("hw.packages")]
@@ -20,6 +20,7 @@ module Gigatoken
20
20
  option :pretokenizer, desc: "pretokenizer scheme, required when TOKENIZER is a .tiktoken file (one of #{Native.pretokenizer_names.join(", ")}); ignored otherwise"
21
21
 
22
22
  def call(tokenizer:, files:, doc_separator: nil, pretokenizer: nil, **)
23
+ Support.check_usage!(files, doc_separator)
23
24
  gt_tokenizer = Support.load_tokenizer(tokenizer, pretokenizer: pretokenizer)
24
25
 
25
26
  via_files = gt_tokenizer.encode_files(Support.text_file_source(files, doc_separator))
@@ -38,7 +39,7 @@ module Gigatoken
38
39
  end
39
40
 
40
41
  out.puts "validation OK: #{via_files.length} documents match"
41
- rescue Gigatoken::Error => e
42
+ rescue Gigatoken::Error, SystemCallError => e
42
43
  err.puts "error: #{e.message}"
43
44
  exit(1)
44
45
  end
@@ -40,6 +40,10 @@ module Gigatoken
40
40
  HARMONY_RESERVED_TAIL = (200013..201087).to_h { |id| ["<|reserved_#{id}|>", id] }.freeze
41
41
  private_constant :HARMONY_RESERVED_TAIL
42
42
 
43
+ # Deep-frozen: entries, their `special_tokens` tables and the `rank_file`
44
+ # paths. `Tokenizer#special_tokens` hands the registry's own Hash back to
45
+ # callers, so anything less lets one caller's poke rewrite what every
46
+ # later `from_encoding` in the process loads.
43
47
  REGISTRY = {
44
48
  "r50k_base" => {
45
49
  rank_file: File.join(DATA_DIR, "r50k_base.tiktoken"),
@@ -67,7 +71,7 @@ module Gigatoken
67
71
  pretokenizer: "o200k",
68
72
  special_tokens: HARMONY_HEAD_TOKENS.merge(HARMONY_RESERVED_TAIL).freeze
69
73
  }
70
- }.freeze
74
+ }.each_value { |encoding| encoding.each_value(&:freeze).freeze }.freeze
71
75
  private_constant :REGISTRY
72
76
 
73
77
  # The packaged encoding names — the single source error messages naming
@@ -88,15 +92,16 @@ module Gigatoken
88
92
 
89
93
  class << self
90
94
  # The {rank_file:, pretokenizer:, special_tokens:} registered for a
91
- # packaged encoding name, or nil.
95
+ # packaged encoding name, or nil. Names are Strings or Symbols, as
96
+ # Tokenizer.load accepts both.
92
97
  def [](name)
93
- REGISTRY[name]
98
+ REGISTRY[name.to_s]
94
99
  end
95
100
 
96
101
  # Why `name` can't be packaged, or nil when there's no reason on
97
102
  # record (it's either packaged, or simply not one gigatoken knows of).
98
103
  def unpackable_reason(name)
99
- UNPACKABLE_REASONS[name]
104
+ UNPACKABLE_REASONS[name.to_s]
100
105
  end
101
106
  end
102
107
  end
data/lib/gigatoken/hub.rb CHANGED
@@ -2,6 +2,7 @@
2
2
 
3
3
  require "async"
4
4
  require "async/http"
5
+ require "async/http/proxy"
5
6
  require "fileutils"
6
7
  require "pathname"
7
8
 
@@ -32,6 +33,39 @@ module Gigatoken
32
33
  MAX_REDIRECTS = 10
33
34
  private_constant :MAX_REDIRECTS
34
35
 
36
+ # Connect/read timeout in seconds, and the bound on the request phase —
37
+ # huggingface_hub's HF_HUB_ETAG_TIMEOUT / HF_HUB_DOWNLOAD_TIMEOUT
38
+ # default.
39
+ DEFAULT_TIMEOUT = 10
40
+ private_constant :DEFAULT_TIMEOUT
41
+
42
+ # A "." or ".." path segment: the traversal that must never reach a URL
43
+ # or the cache, wherever it comes from.
44
+ DOT_SEGMENT = /\A\.\.?\z/
45
+ private_constant :DOT_SEGMENT
46
+
47
+ # Everything outside RFC 3986's unreserved set, which is what
48
+ # huggingface_hub's `quote` percent-encodes.
49
+ RESERVED = /[^A-Za-z0-9\-._~]/
50
+ private_constant :RESERVED
51
+
52
+ # How a request, or the body read that follows it, fails underneath
53
+ # async-http: a connect/read timeout, a refused or reset connection, a
54
+ # proxy refusing the CONNECT tunnel, DNS resolution, a body cut short, a
55
+ # malformed endpoint, TLS. All of them reach the caller as
56
+ # Gigatoken::HubError.
57
+ TRANSPORT_ERRORS = [
58
+ Async::TimeoutError,
59
+ Async::HTTP::Proxy::ConnectFailure,
60
+ IOError,
61
+ SocketError,
62
+ SystemCallError,
63
+ URI::InvalidURIError,
64
+ Protocol::HTTP::Error,
65
+ OpenSSL::SSL::SSLError
66
+ ].freeze
67
+ private_constant :TRANSPORT_ERRORS
68
+
35
69
  class << self
36
70
  # Whether `name` is shaped like a HuggingFace Hub repo id: `org/name`,
37
71
  # or a bare legacy repo name like `gpt2`. At most one slash, and not
@@ -105,15 +139,44 @@ module Gigatoken
105
139
  env("HF_ENDPOINT") || DEFAULT_ENDPOINT
106
140
  end
107
141
 
142
+ # Whether `value` is usable as a URL path and a cache path component:
143
+ # not empty, not absolute, no NUL byte, no "." or ".." segment.
144
+ # huggingface_hub rejects the same traversals in validate_repo_id.
145
+ def safe_component?(value)
146
+ !value.empty? && !value.include?("\0") && !value.start_with?("/") &&
147
+ value.split("/", -1).none? { |segment| segment.match?(DOT_SEGMENT) }
148
+ end
149
+
150
+ # The proxy URL the environment names for a `scheme` request to
151
+ # `hostname`, or nil: http_proxy/https_proxy, lowercase spelling
152
+ # first, suppressed by a matching no_proxy entry — what requests does,
153
+ # and so what huggingface_hub inherits.
154
+ def proxy_url(scheme, hostname)
155
+ url = env("#{scheme}_proxy") || env("#{scheme.upcase}_PROXY")
156
+ url unless url.nil? || no_proxy?(hostname)
157
+ end
158
+
108
159
  private
109
160
 
161
+ # no_proxy is a comma-separated list of host suffixes, or "*" for
162
+ # everything.
163
+ def no_proxy?(hostname)
164
+ host = hostname.downcase
165
+ (env("no_proxy") || env("NO_PROXY")).to_s.split(",").any? do |entry|
166
+ entry = entry.strip.downcase.delete_prefix(".")
167
+ next true if entry == "*"
168
+
169
+ !entry.empty? && (host == entry || host.end_with?(".#{entry}"))
170
+ end
171
+ end
172
+
110
173
  def env(key)
111
174
  value = ENV[key]
112
175
  value unless value.nil? || value.empty?
113
176
  end
114
177
 
115
178
  def word_part?(part, first_alnum:)
116
- return false if part.nil? || part.empty?
179
+ return false if part.nil? || part.empty? || part.match?(DOT_SEGMENT)
117
180
 
118
181
  first_ok = first_alnum ? part[0].match?(/[A-Za-z0-9]/) : word_char?(part[0])
119
182
  first_ok && part[1..].chars.all? { |c| word_char?(c) }
@@ -133,14 +196,26 @@ module Gigatoken
133
196
  # for pointing at a local server in tests (dependency injection, not a
134
197
  # mock); defaults to Hub.default_endpoint (HF_ENDPOINT, then
135
198
  # huggingface.co).
136
- def initialize(endpoint: self.class.default_endpoint)
199
+ # @parameter timeout [Numeric] connect/read timeout in seconds, applied
200
+ # to every request and to the request phase as a whole.
201
+ def initialize(endpoint: self.class.default_endpoint, timeout: DEFAULT_TIMEOUT)
137
202
  @endpoint = endpoint.chomp("/")
138
- @internet = Async::HTTP::Internet.new
203
+ @timeout = timeout
139
204
  end
140
205
 
141
206
  # Path of `filename` from Hub repo `repo_id` at `revision`, served from
142
- # the standard HF cache, downloading into it first when absent.
207
+ # the standard HF cache, downloading into it first when absent. The
208
+ # three caller-supplied components are checked before any request or
209
+ # filesystem access: one of them carrying `..` would otherwise read and
210
+ # overwrite files outside the cache.
143
211
  def hub_file(repo_id, filename = "tokenizer.json", revision: "main")
212
+ {"repo id" => repo_id, "filename" => filename, "revision" => revision}.each do |what, value|
213
+ next if self.class.safe_component?(value)
214
+
215
+ raise HubError, "#{what} #{value.inspect}: must not be empty or absolute, " \
216
+ "or contain a NUL byte or a \".\" or \"..\" path segment"
217
+ end
218
+
144
219
  self.class.cached_file(repo_id, filename, revision) ||
145
220
  Sync { fetch(repo_id, filename, revision) }
146
221
  end
@@ -152,31 +227,127 @@ module Gigatoken
152
227
  # recording the branch ref so later lookups (ours and
153
228
  # huggingface_hub's) resolve it.
154
229
  def fetch(repo_id, filename, revision)
155
- url = "#{@endpoint}/#{repo_id}/resolve/#{revision}/#{filename}"
230
+ clients = []
231
+ url = resolve_url(repo_id, filename, revision)
156
232
  token = self.class.hf_token
157
- response = @internet.get(url, auth_headers(token))
233
+ headers = auth_headers(token)
234
+ response = get(url, headers, clients)
158
235
  # Unlisted headers parse as a Header::Generic (an Array of values);
159
236
  # x-repo-commit is always a single value, so flatten it to a String.
160
237
  commit = response.headers["x-repo-commit"]&.to_s
161
238
 
162
- # Redirects are followed by hand: resolve/ URLs answer with the
163
- # x-repo-commit header and a redirect to a CDN for LFS files, and the
164
- # Authorization header must not travel to the other host.
239
+ # Redirects are followed by hand: a resolve/ URL answers an LFS file
240
+ # with a redirect to a CDN, a renamed repo with one to its new name,
241
+ # and both headers below need a rule of their own across the hop.
165
242
  hops = 0
166
243
  while (300...400).cover?(response.status)
167
- ensure_ok!(url, response.status, !!token)
168
- hops += 1
169
- raise Error, "#{url}: too many redirects" if hops > MAX_REDIRECTS
170
-
171
- location = response.headers["location"] ||
172
- raise(Error, "#{url}: redirect with no Location header")
244
+ location = response.headers["location"]
173
245
  response.close
174
- url = absolutize(location, url)
175
- response = @internet.get(url, {"user-agent" => "gigatoken"})
246
+ raise HubError, "#{url}: redirect with no Location header" unless location
247
+ raise HubError, "#{url}: too many redirects" if (hops += 1) > MAX_REDIRECTS
248
+
249
+ target = absolutize(location, url)
250
+ # A renamed repo answers with a same-origin redirect that still needs
251
+ # the token; the LFS CDN elsewhere authenticates by signed URL and
252
+ # must never see it — huggingface_hub drops the header on the same
253
+ # rule.
254
+ headers = auth_headers(nil) unless same_origin?(target, url)
255
+ url = target
256
+ response = get(url, headers, clients)
257
+ # An LFS file carries x-repo-commit on the resolve/ hop; a renamed
258
+ # repo carries it only on the hop that finally answers 200.
259
+ commit ||= response.headers["x-repo-commit"]&.to_s
176
260
  end
177
- ensure_ok!(url, response.status, !!token)
261
+ ensure_ok!(url, response, !!token)
262
+ ensure_commit!(url, commit)
263
+
264
+ write_to_cache(repo_id, filename, revision, commit, response)
265
+ rescue *TRANSPORT_ERRORS => e
266
+ raise HubError, "#{url}: #{e.message} (#{e.class})"
267
+ ensure
268
+ close_all(response, clients)
269
+ end
178
270
 
179
- write_to_cache(repo_id, filename, revision, commit || revision, response)
271
+ # Release everything the fetch holds, in the one place every exit passes
272
+ # through. Order matters: an unread body keeps its connection checked
273
+ # out, and Async::HTTP::Client#close waits for its pool to drain, so the
274
+ # response goes first and the clients innermost-first — a tunnel client's
275
+ # connection is held by the proxy client opened under it. Bounded, and
276
+ # deaf to the ways a broken connection fails to close: a connection some
277
+ # failure left mid-stream would otherwise stall the drain, and with it
278
+ # the Sync this all runs inside, for good. A leaked socket is worth less
279
+ # than the caller's result or exception.
280
+ def close_all(response, clients)
281
+ Async::Task.current.with_timeout(@timeout) do
282
+ response&.close
283
+ clients.reverse_each(&:close)
284
+ end
285
+ rescue *TRANSPORT_ERRORS
286
+ nil
287
+ end
288
+
289
+ # `endpoint/repo/resolve/revision/filename`, percent-encoded the way
290
+ # huggingface_hub's `quote` does it: `revision` whole (`safe=""`), so
291
+ # `refs/pr/1` travels as `refs%2Fpr%2F1`, while the repo id and the
292
+ # filename keep their slashes.
293
+ def resolve_url(repo_id, filename, revision)
294
+ "#{@endpoint}/#{escape_path(repo_id)}/resolve/#{escape(revision)}/#{escape_path(filename)}"
295
+ end
296
+
297
+ def escape_path(value)
298
+ value.split("/", -1).map { |segment| escape(segment) }.join("/")
299
+ end
300
+
301
+ def escape(value)
302
+ value.gsub(RESERVED) { |char| char.bytes.map { |byte| format("%%%02X", byte) }.join }
303
+ end
304
+
305
+ # GET `url`, through the proxy the environment names for it when there
306
+ # is one: an http request travels to the proxy with the absolute URI in
307
+ # its request line, an https one through a CONNECT tunnel — how requests,
308
+ # and so huggingface_hub, routes them. The clients stay open (the body
309
+ # is still to be streamed); #fetch closes them.
310
+ def get(url, headers, clients)
311
+ endpoint = endpoint_for(url)
312
+ proxy = self.class.proxy_url(endpoint.scheme, endpoint.hostname)
313
+ proxy &&= endpoint_for(proxy)
314
+
315
+ client =
316
+ if proxy.nil?
317
+ open_client(endpoint, clients)
318
+ elsif endpoint.secure?
319
+ # https tunnels through the proxy with CONNECT, then speaks TLS to
320
+ # the origin as if it had connected to it directly.
321
+ open_client(open_client(proxy, clients).proxied_endpoint(endpoint), clients)
322
+ else
323
+ open_client(proxy, clients)
324
+ end
325
+ # An http proxy is addressed with the absolute URI in the request line;
326
+ # a direct or tunnelled request carries just the path.
327
+ target = (proxy && !endpoint.secure?) ? url : endpoint.path
328
+
329
+ request = Protocol::HTTP::Request["GET", target, headers, scheme: endpoint.scheme, authority: endpoint.authority]
330
+ # A CONNECT tunnel's socket carries no timeout of its own, so the
331
+ # request phase is bounded here whichever way it was routed; the body
332
+ # then streams under the endpoint's own connect/read timeout.
333
+ Async::Task.current.with_timeout(@timeout) { client.call(request) }
334
+ end
335
+
336
+ # A malformed URL comes back as URI::InvalidURIError, one that cannot be
337
+ # routed (no scheme or host) as ArgumentError; both mean the same thing
338
+ # to the caller, and neither is worth a backtrace.
339
+ def endpoint_for(url)
340
+ Async::HTTP::Endpoint.parse(url, timeout: @timeout)
341
+ rescue ArgumentError => e
342
+ raise HubError, "#{url}: #{e.message} (#{e.class})"
343
+ end
344
+
345
+ # A client for one hop, remembered in `clients` so #fetch can close it
346
+ # once the body is written. retries: 1 — a client built for this hop has
347
+ # no stale pooled connection for a retry to rescue, and async-http's
348
+ # default of 3 would pay @timeout three times over.
349
+ def open_client(endpoint, clients)
350
+ Async::HTTP::Client.new(endpoint, retries: 1).tap { |client| clients << client }
180
351
  end
181
352
 
182
353
  def auth_headers(token)
@@ -199,51 +370,75 @@ module Gigatoken
199
370
  rescue
200
371
  FileUtils.rm_f(tmp)
201
372
  raise
202
- ensure
203
- response.close
204
373
  end
205
374
  File.rename(tmp, target)
206
375
 
207
- if !self.class.commit_hash?(revision) && revision != commit
208
- refs_dir = repo_dir.join("refs")
209
- FileUtils.mkdir_p(refs_dir)
210
- File.write(refs_dir.join(revision), commit)
376
+ # A nested revision like refs/pr/1 is a nested ref file, so its parent
377
+ # has to exist — huggingface_hub mkdir -p's the same path.
378
+ if revision != commit
379
+ ref_path = repo_dir.join("refs", revision)
380
+ FileUtils.mkdir_p(ref_path.dirname)
381
+ File.write(ref_path, commit)
211
382
  end
212
383
 
213
384
  target
214
385
  end
215
386
 
216
- def ensure_ok!(url, status, had_token)
387
+ # Raise on a non-success status. The unread body is left to #close_all,
388
+ # which every exit from #fetch passes through: it keeps the HTTP/1.x
389
+ # connection — and the Sync around it — alive until it is closed, so the
390
+ # exception would otherwise never reach the caller.
391
+ def ensure_ok!(url, response, had_token)
392
+ status = response.status
217
393
  return if (200...400).cover?(status)
218
394
 
219
395
  case status
220
396
  when 404
221
- raise Error, "#{url}: HTTP 404 — no such repo with that file, and no such local file either"
397
+ raise HubError, "#{url}: HTTP 404 — no such repo with that file, and no such local file either"
222
398
  when 401, 403
223
399
  token_note = had_token ? "the request used the discovered token" : "no token was found"
224
- raise Error,
400
+ raise HubError,
225
401
  "#{url}: HTTP #{status} — the repo may be private or gated (#{token_note}; set HF_TOKEN or run " \
226
402
  "`hf auth login`, and accept the repo's terms on huggingface.co if it is gated)"
227
403
  else
228
- raise Error, "#{url}: HTTP #{status}"
404
+ raise HubError, "#{url}: HTTP #{status}"
229
405
  end
230
406
  end
231
407
 
408
+ # The snapshot directory is named by a server-controlled header, so only
409
+ # a real commit hash may become one: anything else would write the body
410
+ # wherever the header points. A missing header also means the endpoint
411
+ # isn't a Hub, whose snapshot would never be found again.
412
+ def ensure_commit!(url, commit)
413
+ return if self.class.commit_hash?(commit.to_s)
414
+
415
+ raise HubError, "#{url}: response is missing a usable x-repo-commit header (#{commit.inspect}) — it does not " \
416
+ "seem to be served by a HuggingFace Hub endpoint; if HF_ENDPOINT is set, check that it points to a " \
417
+ "Hub-compatible endpoint, and otherwise check your firewall and proxy settings"
418
+ end
419
+
232
420
  # A redirect Location resolved against the request URL: absolute URLs
233
421
  # pass through, host-relative (`/x/y`) and path-relative ones join the
234
422
  # base.
235
423
  def absolutize(location, base)
236
424
  return location if location.include?("://")
237
425
 
238
- origin_end = base.index("://") ? base.index("://") + 3 : 0
239
- origin_end = base.index("/", origin_end) || base.length
426
+ base_origin = origin(base)
427
+ return "#{base_origin}#{location}" if location.start_with?("/")
240
428
 
241
- if location.start_with?("/")
242
- "#{base[0...origin_end]}#{location}"
243
- else
244
- dir_end = base.rindex("/") || base.length
245
- "#{base[0...[dir_end, origin_end].max]}/#{location}"
246
- end
429
+ dir_end = base.rindex("/") || base.length
430
+ "#{base[0...[dir_end, base_origin.length].max]}/#{location}"
431
+ end
432
+
433
+ # Scheme and authority of a URL — everything before its path. Two URLs
434
+ # sharing one may pass the Authorization header between them.
435
+ def origin(url)
436
+ scheme_end = url.index("://")
437
+ url[0...(url.index("/", scheme_end ? scheme_end + 3 : 0) || url.length)]
438
+ end
439
+
440
+ def same_origin?(url, other)
441
+ origin(url).casecmp?(origin(other))
247
442
  end
248
443
  end
249
444
  end
@@ -29,8 +29,12 @@ module Gigatoken
29
29
  end
30
30
 
31
31
  # Array of token ids for document `i`, materialized on demand; negative
32
- # indices count from the end and out-of-range ones give nil, like Array.
33
- def [](i)
32
+ # indices count from the end, out-of-range ones give nil, and a
33
+ # non-Integer index is a TypeError, all like Array.
34
+ def [](index)
35
+ i = Integer.try_convert(index)
36
+ raise TypeError, "no implicit conversion of #{index.class} into Integer" if i.nil?
37
+
34
38
  i += size if i.negative?
35
39
  return unless i >= 0 && i < size
36
40