Class: EmbeddingUtil::CLI

Inherits:
Thor
  • Object
show all
Defined in:
lib/embedding_util/cli.rb

Constant Summary collapse

CONFIG_OPTIONS = 
{
  profile: :to_sym.to_proc,
  runtime: ->(value) { RuntimeCommand.normalize_runtime(value) },
  endpoint: ->(value) { value },
  embedding_endpoint: ->(value) { value },
  reranker_endpoint: ->(value) { value },
  timeout: ->(value) { value },
  startup_timeout: ->(value) { value },
  shutdown_idle: :to_i.to_proc,
  reranker_ubatch_size: :to_i.to_proc,
  reranker_max_ubatch_size: :to_i.to_proc,
  verbose: ->(value) { value }
}.freeze

Instance Method Summary collapse

Instance Method Details

#embed(text) ⇒ Object 



55
56
57
58
59
60
 
# File 'lib/embedding_util/cli.rb', line 55

def embed(text)
  configure_embedding_util
  puts JSON.generate(EmbeddingUtil.embed(text))
rescue Error => e
  abort e.message
end

#profiles ⇒ Object 



46
47
48
49
50
51
52
 
# File 'lib/embedding_util/cli.rb', line 46

def profiles
  EmbeddingUtil.profiles.each do |profile|
    puts profile.name
    puts "  embedding: #{profile.embedding.fetch(:repo)} / #{profile.embedding.fetch(:file)}"
    puts "  reranker:  #{profile.reranker.fetch(:repo)} / #{profile.reranker.fetch(:file)}"
  end
end

#rerank(query, *documents) ⇒ Object 



63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
 
# File 'lib/embedding_util/cli.rb', line 63

def rerank(query, *documents)
  configure_embedding_util
  raise Error, "provide at least one document to rerank" if documents.empty?

  results = EmbeddingUtil.rerank(query, documents).map do |result|
    {
      index: result.index,
      document: result.document,
      score: result.score,
      metadata: result.metadata
    }
  end
  puts JSON.pretty_generate(results)
rescue Error => e
  abort e.message
end

#serve ⇒ Object 



85
86
87
88
89
90
91
92
93
94
95
96
97
98
 
# File 'lib/embedding_util/cli.rb', line 85

def serve
  configure_embedding_util
  ServerManager.new(config: EmbeddingUtil.configuration).serve(
    model: options[:model],
    runtime: options[:runtime] || EmbeddingUtil.configuration.runtime,
    shutdown_idle: options[:shutdown_idle]&.to_i,
    host: options[:host],
    port: options[:port]&.to_i
  )
rescue Error => e
  abort e.message
rescue Interrupt
  exit 130
end

#support ⇒ Object 



35
36
37
38
39
40
41
42
43
 
# File 'lib/embedding_util/cli.rb', line 35

def support
  configure_embedding_util
  EmbeddingUtil.support.each do |item|
    status = item.fetch(:supported) ? "supported" : "not supported"
    puts "#{item.fetch(:provider)}: #{status}"
    puts "  embedding_endpoint: #{item.fetch(:embedding_endpoint)}" if item[:embedding_endpoint]
    puts "  reranker_endpoint: #{item.fetch(:reranker_endpoint)}" if item[:reranker_endpoint]
  end
end