Kohagi (Ruby)

A Ruby client for Kohagi, a local CLI for generating sentence embeddings.

It runs the kohagi binary and safely handles its input and output streams. It also reports execution errors and returns embeddings associated with their input IDs.

This gem handles only communication with the CLI, converting {id, text} records into {id, embedding} results. How the embeddings are stored, validated, or used is left to your application.

Requirements

The kohagi binary must be available on PATH, or specified with the bin: option.

See the Kohagi installation instructions.

This gem ships no native code and does not bundle the kohagi binary.

Install

gem "kohagi"

Usage

require "kohagi"

client = Kohagi::Client.new(prefix: "検索文書: ")

records = [
  { id: 1, text: "夏目漱石『吾輩は猫である』…" },
  { id: 2, text: "太宰治『人間失格』…" },
]

# Streaming: each result is yielded as it arrives, so memory stays flat on any corpus. Returns a Summary.
summary = client.embed(records) do |r|
  store(r.id, r.embedding)   # e.g. write to a pgvector column
end

summary.dim        # => 512
summary.out        # => 2   (records embedded)
summary.skipped    # => 0
summary.truncated  # => 0   (records that ran past --max-seq-length)

Without a block the results are collected onto the summary:

summary = client.embed(records)
summary.results.each { |r| store(r.id, r.embedding) }

Configuration

Each option maps directly to a kohagi CLI flag. Kohagi does not use a configuration file.

Option	CLI flag or behavior
bin:	Executable path. Defaults to "kohagi".
model_id:	--model-id
model_path: and tokenizer_path:	--model-path and --tokenizer-path for offline use. These take precedence over model_id:.
prefix:	--prefix
pooling:	--pooling
device:	--device ("cpu", "metal", or "coreml")
coreml_dir:, coreml_model_id:, and coreml_prefer:	Corresponding --coreml-* flags
max_seq_length:, batch_size:, and precision:	Corresponding CLI flags
normalize:	Passes --no-normalize when false. Normalization is enabled by default.
report_tokens:	Passes --report-tokens, adding n_tokens and truncated? to each result.
logger:	A callable that receives kohagi's standard error output after each run.

Truncation

Text longer than max_seq_length is truncated before embedding. The default limit is 512 tokens.

The summary always includes the number of truncated inputs. With report_tokens: true, each result also includes n_tokens and truncated?. You can use this information to process truncated inputs separately:

client = Kohagi::Client.new(report_tokens: true)

client.embed(records) do |result|
  chunk_and_reembed(result.id) if result.truncated?
end

Exit codes

kohagi exit code	Result
0	Returns a Summary for which summary.ok? is true.
2	Returns a Summary for which summary.partial? is true. Some input lines were skipped, but the returned output is valid.
1	Raises Kohagi::FatalError.
3	Raises Kohagi::UnsupportedDeviceError.

Exit code 3 indicates that the Core ML backend cannot handle the request. You can retry using the CPU:

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

License

MIT