Class: StaticEmbeddings::Reference

Inherits:
Object
  • Object
show all
Defined in:
lib/static_embeddings/reference.rb

Constant Summary collapse

CJK_RANGES = 
[
  0x4E00..0x9FFF, 0x3400..0x4DBF, 0x20000..0x2A6DF, 0x2A700..0x2B73F,
  0x2B740..0x2B81F, 0x2B820..0x2CEAF, 0xF900..0xFAFF, 0x2F800..0x2FA1F
].freeze
ASCII_PUNCT = 
[33..47, 58..64, 91..96, 123..126].freeze
ASCII_SPACES = 
[" ", "\t", "\n", "\r"].freeze
RE_PUNCT = 
/\A\p{P}\z/
RE_CONTROL = 
/\A(?:\p{Cc}|\p{Cf}|\p{Co}|\p{Cs})\z/
RE_MN = 
/\A\p{Mn}\z/
RE_WHITESPACE = 
/\A(?:\p{Zs}|[\u0085\u2028\u2029])\z/

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(tokens:, matrix:, meta:) ⇒ Reference

Returns a new instance of Reference.



17
18
19
20
21
22
23
 
# File 'lib/static_embeddings/reference.rb', line 17

def initialize(tokens:, matrix:, meta:)
  @tokens = tokens
  @vocab = tokens.each_with_index.to_h
  @matrix = matrix
  @meta = meta
  @dim = meta.fetch(:dim)
end

Instance Attribute Details

#meta ⇒ Object (readonly)

Returns the value of attribute meta.



15
16
17
 
# File 'lib/static_embeddings/reference.rb', line 15

def meta
  @meta
end

Class Method Details

.from_source_dir(dir, max_tokens: Converter::REFERENCE_MAX_TOKENS) ⇒ Object 



25
26
27
28
29
30
31
32
33
34
35
36
37
38
 
# File 'lib/static_embeddings/reference.rb', line 25

def self.from_source_dir(dir, max_tokens: Converter::REFERENCE_MAX_TOKENS)
  tokenizer = JSON.parse(File.binread(File.join(dir, "tokenizer.json")))
  config_path = File.join(dir, "config.json")
  config = File.file?(config_path) ? JSON.parse(File.binread(config_path)) : {}
  tokens = tokens_from(tokenizer)
  tensor = Safetensors.read(File.join(dir, "model.safetensors"))[:tensors].values.first
  rows, dim = tensor[:shape]
  floats = tensor[:bytes].unpack("e*")
  normalizer = tokenizer["normalizer"] || {}

  new(tokens: tokens,
      matrix: Array.new(rows) { |i| floats[i * dim, dim] },
      meta: meta_from(tokenizer, config, normalizer, dim, max_tokens))
end

.meta_from(tokenizer, config, normalizer, dim, max_tokens) ⇒ Object 



45
46
47
48
49
50
51
52
53
54
55
56
57
58
 
# File 'lib/static_embeddings/reference.rb', line 45

def self.meta_from(tokenizer, config, normalizer, dim, max_tokens)
  lowercase = normalizer.fetch("lowercase", true)
  {
    dim: dim,
    lowercase: lowercase,
    strip_accents: normalizer["strip_accents"].nil? ? lowercase : normalizer["strip_accents"],
    clean_text: normalizer.fetch("clean_text", true),
    handle_chinese_chars: normalizer.fetch("handle_chinese_chars", true),
    max_input_chars_per_word: tokenizer.dig("model", "max_input_chars_per_word") || 100,
    unk_token: tokenizer.dig("model", "unk_token") || "[UNK]",
    normalize: config.key?("normalize") ? config["normalize"] : true,
    max_tokens: max_tokens
  }
end

.tokens_from(tokenizer) ⇒ Object 



40
41
42
43
 
# File 'lib/static_embeddings/reference.rb', line 40

def self.tokens_from(tokenizer)
  vocab = tokenizer.dig("model", "vocab")
  vocab.each_with_object(Array.new(vocab.length)) { |(token, id), tokens| tokens[id] = token }
end

Instance Method Details

#embed(text, max_tokens: ) ⇒ Object 



79
80
81
82
83
84
85
 
# File 'lib/static_embeddings/reference.rb', line 79

def embed(text, max_tokens: @meta[:max_tokens])
  used = tokenize(text, max_tokens: max_tokens).reject { |id| id == unk_id }
  return Array.new(@dim, 0.0) if used.empty?

  vector = pooled(used)
  @meta[:normalize] ? l2_normalize(vector) : vector
end

#normalize_text(text) ⇒ Object 



60
61
62
63
64
65
66
67
 
# File 'lib/static_embeddings/reference.rb', line 60

def normalize_text(text)
  chars = text.chars
  chars = clean_chars(chars) if @meta[:clean_text]
  chars = split_chinese(chars) if @meta[:handle_chinese_chars]
  chars = strip_accents(chars) if @meta[:strip_accents]
  chars = lowercase(chars) if @meta[:lowercase]
  chars.join
end

#pre_tokenize(normalized) ⇒ Object 



69
70
71
 
# File 'lib/static_embeddings/reference.rb', line 69

def pre_tokenize(normalized)
  normalized.split(" ").flat_map { |word| split_punctuation(word) }
end

#tokenize(text, max_tokens: ) ⇒ Object 



73
74
75
76
77
 
# File 'lib/static_embeddings/reference.rb', line 73

def tokenize(text, max_tokens: @meta[:max_tokens])
  ids = []
  pre_tokenize(normalize_text(text)).each { |word| ids.concat(wordpiece(word)) }
  max_tokens && max_tokens.positive? && ids.length > max_tokens ? ids.first(max_tokens) : ids
end