Class: Cohere::Transcribe::Safetensors::Reader

Inherits:
Object
  • Object
show all
Defined in:
lib/cohere/transcribe/safetensors.rb

Overview

Reads and validates one safetensors file. Tensor payloads stay on disk.

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(path) ⇒ Reader

Returns a new instance of Reader.



38
39
40
41
 
# File 'lib/cohere/transcribe/safetensors.rb', line 38

def initialize(path)
  @path = Pathname(path).expand_path
  @tensors = read_header.freeze
end

Instance Attribute Details

#path ⇒ Object (readonly)

Returns the value of attribute path.



36
37
38
 
# File 'lib/cohere/transcribe/safetensors.rb', line 36

def path
  @path
end

#tensors ⇒ Object (readonly)

Returns the value of attribute tensors.



36
37
38
 
# File 'lib/cohere/transcribe/safetensors.rb', line 36

def tensors
  @tensors
end

Instance Method Details

#fetch(name) ⇒ Object 



51
52
53
54
55
 
# File 'lib/cohere/transcribe/safetensors.rb', line 51

def fetch(name)
  tensors.fetch(name)
rescue KeyError
  raise Error, "Tensor #{name.inspect} is not present in #{path}"
end

#key?(name) ⇒ Boolean

Returns:

  • (Boolean) 


47
48
49
 
# File 'lib/cohere/transcribe/safetensors.rb', line 47

def key?(name)
  tensors.key?(name)
end

#names ⇒ Object 



43
44
45
 
# File 'lib/cohere/transcribe/safetensors.rb', line 43

def names
  tensors.keys
end

#write_tensor(tensor, output, target_dtype:, converter: DTypeConverter.default, chunk_bytes: DEFAULT_CHUNK_BYTES) ⇒ Object

Raises:

  • (ArgumentError) 


57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
 
# File 'lib/cohere/transcribe/safetensors.rb', line 57

def write_tensor(tensor, output, target_dtype:, converter: DTypeConverter.default,
                 chunk_bytes: DEFAULT_CHUNK_BYTES)
  raise ArgumentError, "Tensor #{tensor.name.inspect} belongs to a different safetensors reader" unless tensor.reader.equal?(self)

  target_dtype = target_dtype.to_s.upcase
  unless FLOAT_DTYPES.include?(tensor.dtype) && FLOAT_DTYPES.include?(target_dtype)
    raise Error, "Cannot convert #{tensor.name.inspect} from #{tensor.dtype} to #{target_dtype}"
  end
  raise ArgumentError, "chunk_bytes must be positive" unless chunk_bytes.positive?

  source_width = DTYPE_BYTES.fetch(tensor.dtype)
  elements_per_chunk = [chunk_bytes / source_width, 1].max
  bytes_per_chunk = elements_per_chunk * source_width
  remaining = tensor.nbytes
  written = 0

  File.open(path, "rb") do |source|
    source.seek(tensor.data_start)
    while remaining.positive?
      requested = [remaining, bytes_per_chunk].min
      chunk = source.read(requested)
      raise Error, "Unexpected end of #{path} while reading #{tensor.name.inspect}" if chunk.nil? || chunk.bytesize != requested

      converted = converter.convert(chunk, from: tensor.dtype, to: target_dtype)
      output.write(converted)
      written += converted.bytesize
      remaining -= requested
    end
  end

  expected = tensor.element_count * DTYPE_BYTES.fetch(target_dtype)
  return written if written == expected

  raise Error,
        "Converted tensor #{tensor.name.inspect} wrote #{written} bytes; expected #{expected}"
end