Class: Kohagi::Client

Inherits:
Object
  • Object
show all
Defined in:
lib/kohagi/client.rb

Overview

Drives the kohagi binary's stdin/stdout JSONL protocol.

kohagi is config-file-less, so every option here maps 1:1 to a CLI flag. The client owns only the transport: building the command, spawning the process without deadlocking on the pipe buffer, mapping exit codes to outcomes, and parsing the records back. The domain (which texts, what to do with the vectors) stays with the caller.

client = Kohagi::Client.new(prefix: "検索文書: ", report_tokens: true)
summary = client.embed(records) { |r| store(r.id, r.embedding) }

Constant Summary collapse

DEFAULTS = 
{ bin: "kohagi", normalize: true, report_tokens: false }.freeze

Instance Method Summary collapse

Constructor Details

#initialize(**options) ⇒ Client

See the README for the full option list. Unknown keys are ignored, so a newer kohagi flag can be threaded through as extra_flags.



23
24
25
 
# File 'lib/kohagi/client.rb', line 23

def initialize(**options)
  @options = DEFAULTS.merge(options)
end

Instance Method Details

#command ⇒ Object

The argv kohagi is launched with. Pure and public, so a caller can assert on it without spawning a process.



37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
 
# File 'lib/kohagi/client.rb', line 37

def command
  o = @options
  cmd = [o.fetch(:bin)]
  cmd.concat(model_flags(o))
  cmd.concat(["--prefix", o[:prefix]]) if o[:prefix]
  cmd.concat(device_flags(o))
  {
    "--pooling" => o[:pooling], "--max-seq-length" => o[:max_seq_length],
    "--batch-size" => o[:batch_size], "--precision" => o[:precision]
  }.each { |flag, val| cmd.concat([flag, val.to_s]) unless val.nil? }
  cmd << "--no-normalize" unless o.fetch(:normalize)
  cmd << "--report-tokens" if o.fetch(:report_tokens)
  cmd.concat(Array(o[:extra_flags]))
  cmd
end

#embed(records, &block) ⇒ Object

Embed records (each {id:, text:}). With a block, yields a Result per record as it arrives, so memory stays flat on any corpus; without one, the results are collected into the returned Summary. Returns a Summary either way.

Exit codes become outcomes: 0 and 2 return normally (2 = some lines were skipped, but the output you received is valid); 1 raises FatalError; 3 raises UnsupportedDeviceError so the caller can retry on with(device: "cpu").



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
93
94
95
96
97
98
99
100
101
102
 
# File 'lib/kohagi/client.rb', line 62

def embed(records, &block)
  collected = block ? nil : []
  emit = block || ->(r) { collected << r }
  dim = nil
  out = 0

  Open3.popen3(*command) do |stdin, stdout, stderr, wait|
    # Write from one thread and read from this one — writing the whole
    # corpus before reading anything can deadlock both processes once the
    # pipe buffer fills. Drain stderr concurrently for the same reason.
    writer = Thread.new do
      records.each { |rec| stdin.puts(JSON.generate(id: rec.fetch(:id), text: rec.fetch(:text))) }
    rescue Errno::EPIPE
      # kohagi exited before reading all input (e.g. exit 3, detected up
      # front). The process exit code carries the real reason; the broken
      # pipe here is a symptom, not the error to surface.
    ensure
      begin
        stdin.close
      rescue Errno::EPIPE
        nil
      end
    end
    errors = Thread.new { stderr.read }

    stdout.each_line do |line|
      result = Result.from_json(line)
      dim ||= result.embedding&.size
      out += 1
      emit.call(result)
    end

    writer.join
    stderr_text = errors.value
    @options[:logger]&.call(stderr_text) unless stderr_text.to_s.empty?
    code = wait.value.exitstatus
    raise_for(code, stderr_text)

    return build_summary(stderr_text, sent: records.size, out: out, dim: dim, code: code, results: collected)
  end
end

#with(**overrides) ⇒ Object

A copy with some options overridden — for the exit-3 → CPU fallback:

rescue Kohagi::UnsupportedDeviceError
client.with(device: "cpu").embed(records, &block)


31
32
33
 
# File 'lib/kohagi/client.rb', line 31

def with(**overrides)
  self.class.new(**@options.merge(overrides))
end