kohagi 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 4beb0ae3b0422e028e52f1e04c90e414f740bf412ef4a871d928f7e760730861
+  data.tar.gz: 7ac174c73b676551ba7ce105e396d1b133e298663d00193467520bb0a162d82e
+SHA512:
+  metadata.gz: 2262d1a3cbb528dbe35615e1e548a7d4604c08635d3ea3a0211187c8d1fc12a1e940d5a80d2535a403d8e3b5fdfa071382ab2757efe2ce0c6a1921330d3adb8d
+  data.tar.gz: 312e26729165f64e008efaa94b83aff99430206ead9219a8a8d2e0fae67505ae385a93fe084cb42ac39107680b71d59cbcdb4a5d9358cba634a535073cc758bf

data/CHANGELOG.md

@@ -0,0 +1,15 @@
+# Changelog
+
+## [0.1.0] - 2026-07-24
+
+Initial release.
+
+- `Kohagi::Client` — builds the kohagi command from 1:1 CLI-flag options,
+  spawns it deadlock-safely (writer thread + concurrent stderr drain), and
+  streams `Kohagi::Result` records back.
+- Exit codes as outcomes: 0/2 return a `Kohagi::Summary` (`ok?` / `partial?`),
+  1 raises `Kohagi::FatalError`, 3 raises `Kohagi::UnsupportedDeviceError` for
+  the CoreML → CPU fallback.
+- `--report-tokens` surfaced as `Result#truncated?` / `#n_tokens` and the
+  summary's `truncated` count.
+- `Client#with` for option overrides (the fallback pattern).

data/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 takahashim
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/README.md

@@ -0,0 +1,112 @@
+# Kohagi (Ruby)
+
+A Ruby client for [Kohagi](https://github.com/takahashim/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](https://github.com/takahashim/kohagi#install).
+
+This gem ships no native code and does not bundle the `kohagi` binary.
+
+## Install
+
+```ruby
+gem "kohagi"
+```
+
+## Usage
+
+```ruby
+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:
+
+```ruby
+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:
+
+```ruby
+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:
+
+```ruby
+begin
+  client.embed(records, &block)
+rescue Kohagi::UnsupportedDeviceError
+  client.with(device: "cpu").embed(records, &block)
+end
+```
+
+## License
+
+MIT

data/lib/kohagi/client.rb

@@ -0,0 +1,156 @@
+# frozen_string_literal: true
+
+require "open3"
+require "json"
+
+module Kohagi
+  # 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) }
+  #
+  class Client
+    DEFAULTS = { bin: "kohagi", normalize: true, report_tokens: false }.freeze
+
+    # See the README for the full option list. Unknown keys are ignored, so a
+    # newer kohagi flag can be threaded through as `extra_flags`.
+    def initialize(**options)
+      @options = DEFAULTS.merge(options)
+    end
+
+    # A copy with some options overridden — for the exit-3 → CPU fallback:
+    #
+    #   rescue Kohagi::UnsupportedDeviceError
+    #     client.with(device: "cpu").embed(records, &block)
+    def with(**overrides)
+      self.class.new(**@options.merge(overrides))
+    end
+
+    # The argv kohagi is launched with. Pure and public, so a caller can assert
+    # on it without spawning a process.
+    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` (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")`.
+    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
+
+    private
+
+    def raise_for(code, stderr)
+      case code
+      when ExitCode::OK, ExitCode::PARTIAL then nil
+      when ExitCode::UNSUPPORTED then raise UnsupportedDeviceError.new(stderr, exit_code: code)
+      else raise FatalError.new(stderr, exit_code: code)
+      end
+    end
+
+    # Combine kohagi's reported counts with what we counted from the stream,
+    # preferring kohagi's numbers and falling back to the stream when a field
+    # (or the whole summary line) is absent.
+    def build_summary(stderr, sent:, out:, dim:, code:, results:)
+      reported = SummaryLine.parse(stderr)
+      Summary.new(
+        model: reported&.model,
+        dim: reported&.dim || dim,
+        sent: sent,
+        out: reported&.out || out,
+        skipped: reported&.skipped || (sent - out),
+        truncated: reported&.truncated,
+        exit_code: code,
+        results: results
+      )
+    end
+
+    # Local weights win over a Hub id; both absent lets kohagi use its default.
+    def model_flags(o)
+      if o[:model_path] && o[:tokenizer_path]
+        ["--model-path", o[:model_path].to_s, "--tokenizer-path", o[:tokenizer_path].to_s]
+      elsif o[:model_id]
+        ["--model-id", o[:model_id].to_s]
+      else
+        []
+      end
+    end
+
+    # --device, plus the converted-model pointer CoreML needs (--coreml-dir
+    # wins over --coreml-model-id, matching kohagi).
+    def device_flags(o)
+      return [] unless o[:device]
+
+      flags = ["--device", o[:device].to_s]
+      return flags unless o[:device].to_s == "coreml"
+
+      flags.concat(["--coreml-dir", o[:coreml_dir].to_s]) if o[:coreml_dir]
+      flags.concat(["--coreml-model-id", o[:coreml_model_id].to_s]) if o[:coreml_model_id] && !o[:coreml_dir]
+      flags.concat(["--coreml-prefer", o[:coreml_prefer].to_s]) if o[:coreml_prefer]
+      flags
+    end
+  end
+end

data/lib/kohagi/errors.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module Kohagi
+  # Base for every error the client raises. Carries the process exit code and
+  # the captured stderr so the caller can log or branch on them.
+  class Error < StandardError
+    attr_reader :exit_code, :stderr
+
+    def initialize(stderr = nil, exit_code: nil)
+      @stderr = stderr
+      @exit_code = exit_code
+      super(build_message(stderr, exit_code))
+    end
+
+    private
+
+    def build_message(stderr, code)
+      base = "kohagi exited #{code}"
+      tail = stderr.to_s.lines.last(3).map(&:chomp).reject(&:empty?).join(" / ")
+      tail.empty? ? base : "#{base}: #{tail}"
+    end
+  end
+
+  # Exit 1: model load failure, bad flags, I/O error. Nothing to retry.
+  class FatalError < Error; end
+
+  # Exit 3: the requested CoreML backend cannot serve the request (built without
+  # the feature, no converted model given, or a sequence past the largest
+  # bucket). Detected before any input is read, so no output was produced — the
+  # caller can retry on `client.with(device: "cpu")`.
+  class UnsupportedDeviceError < Error; end
+end

data/lib/kohagi/exit_code.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Kohagi
+  # The exit codes kohagi uses (see PROTOCOL.md), named in one place so the
+  # raise-vs-return decision (Client#raise_for) and the outcome classification
+  # (Summary#ok? / #partial?) can't drift apart.
+  module ExitCode
+    OK = 0          # every record embedded
+    PARTIAL = 2     # finished, but some input lines were skipped; output valid
+    FATAL = 1       # model load failure, bad flags, I/O error
+    UNSUPPORTED = 3 # the requested CoreML backend cannot serve the request
+  end
+end

data/lib/kohagi/result.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Kohagi
+  # One output record: the echoed id and its vector, plus the optional
+  # `--report-tokens` fields. `id` is whatever JSON value was sent, unchanged.
+  Result = Struct.new(:id, :embedding, :n_tokens, :truncated, keyword_init: true) do
+    # True only when kohagi reported this text as truncated (needs
+    # `report_tokens: true`); nil/false otherwise.
+    def truncated?
+      truncated == true
+    end
+
+    # Parse one JSONL line from kohagi's stdout.
+    def self.from_json(line)
+      row = JSON.parse(line)
+      new(
+        id: row["id"],
+        embedding: row["embedding"],
+        n_tokens: row["n_tokens"],
+        truncated: row["truncated"]
+      )
+    end
+  end
+end

data/lib/kohagi/summary.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Kohagi
+  # The run's outcome: kohagi's reported counts (see SummaryLine) combined by the
+  # client with what actually streamed back. A plain data holder — parsing and
+  # composition live in SummaryLine and Client, not here.
+  #
+  # `truncated` is nil against a kohagi old enough not to report it.
+  Summary = Struct.new(
+    :model, :dim, :sent, :out, :skipped, :truncated, :exit_code, :results,
+    keyword_init: true
+  ) do
+    # Every record embedded.
+    def ok?
+      exit_code == ExitCode::OK
+    end
+
+    # Finished, but some input lines were skipped. The received output is still
+    # valid; investigate stderr and resend the skipped records.
+    def partial?
+      exit_code == ExitCode::PARTIAL
+    end
+  end
+end

data/lib/kohagi/summary_line.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Kohagi
+  # kohagi's one-line run summary from stderr, parsed into a value object:
+  #
+  #   kohagi: model=cl-nagoya/ruri-v3-130m dim=512 in=2141 out=2141 skipped=0 truncated=3
+  #
+  # `SummaryLine.parse` returns an instance, or nil when stderr carries no such
+  # line, so the caller can fall back to values it computed from the stream. A
+  # field the line omits (an older kohagi has no `truncated`) reads as nil.
+  # Knowing kohagi's stderr text format lives here, not in the Summary data.
+  class SummaryLine
+    MARKER = "kohagi: model="
+    MODEL = /\bmodel=(?<value>\S+)/
+    NUMERIC = %i[dim out skipped truncated].freeze
+
+    attr_reader :model, :dim, :out, :skipped, :truncated
+
+    # The parsed summary line in `stderr`, or nil if there isn't one.
+    def self.parse(stderr)
+      line = stderr.to_s.lines.reverse.find { |l| l.include?(MARKER) }
+      line && new(line)
+    end
+
+    def initialize(line)
+      @model = capture(line, MODEL)
+      NUMERIC.each do |name|
+        value = capture(line, /\b#{name}=(?<value>\d+)\b/)
+        instance_variable_set("@#{name}", value&.to_i)
+      end
+    end
+
+    private
+
+    def capture(line, regexp)
+      line.match(regexp)&.[](:value)
+    end
+  end
+end

data/lib/kohagi/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Kohagi
+  VERSION = "0.1.0"
+end

data/lib/kohagi.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require_relative "kohagi/version"
+require_relative "kohagi/exit_code"
+require_relative "kohagi/errors"
+require_relative "kohagi/result"
+require_relative "kohagi/summary_line"
+require_relative "kohagi/summary"
+require_relative "kohagi/client"
+
+# Ruby client for the kohagi local sentence-embedding CLI.
+# See https://github.com/takahashim/kohagi/blob/main/PROTOCOL.md for the
+# stdin/stdout protocol this wraps.
+module Kohagi
+end

metadata

@@ -0,0 +1,58 @@
+--- !ruby/object:Gem::Specification
+name: kohagi
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- takahashim
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies: []
+description: |
+  Drives the kohagi binary's stdin/stdout JSONL protocol from Ruby: builds the
+  command, spawns it without deadlocking on the pipe buffer, maps exit codes to
+  outcomes, and returns id-tagged embeddings. Shells out to an installed kohagi
+  binary and ships no native code of its own.
+email:
+- takahashimm@gmail.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- LICENSE
+- README.md
+- lib/kohagi.rb
+- lib/kohagi/client.rb
+- lib/kohagi/errors.rb
+- lib/kohagi/exit_code.rb
+- lib/kohagi/result.rb
+- lib/kohagi/summary.rb
+- lib/kohagi/summary_line.rb
+- lib/kohagi/version.rb
+homepage: https://github.com/takahashim/kohagi-ruby
+licenses:
+- MIT
+metadata:
+  source_code_uri: https://github.com/takahashim/kohagi-ruby
+  changelog_uri: https://github.com/takahashim/kohagi-ruby/blob/main/CHANGELOG.md
+  rubygems_mfa_required: 'true'
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 4.0.16
+specification_version: 4
+summary: Ruby client for the kohagi local sentence-embedding CLI
+test_files: []