ignis-dl 0.0.1

data/lib/nnw/ai/server.rb

@@ -0,0 +1,268 @@
+# frozen_string_literal: true
+
+require "json"
+require "securerandom" # used for completion IDs; was referenced but never required
+
+module Ignis
+  module AI
+    # OpenAI-compatible inference server.
+    #
+    # Routes:
+    #   POST /v1/completions       — text completion (SSE streaming)
+    #   POST /v1/chat/completions  — chat completion (SSE streaming)
+    #   GET  /v1/models            — list loaded models
+    #   POST /v1/embeddings        — compute embeddings
+    #
+    # Integrates with WNAIS HTTP/2 server if available, otherwise
+    # uses a simple TCP server for standalone operation.
+    class Server
+      # @return [TextGenerator]
+      attr_reader :generator
+
+      # @return [String] model name for API responses
+      attr_reader :model_name
+
+      # @param model [Transformer::Model]
+      # @param tokenizer [Tokenizer]
+      # @param model_name [String]
+      # @param host [String]
+      # @param port [Integer]
+      def initialize(model, tokenizer, model_name: "nnw-model", host: "0.0.0.0", port: 8080)
+        @generator = TextGenerator.new(model, tokenizer)
+        @model_name = model_name
+        @host = host
+        @port = port
+        @batch_processor = BatchProcessor.new(model, tokenizer)
+      end
+
+      # Start the server.
+      # @return [void]
+      def start!
+        Ignis.logger.info("Ignis AI Server starting on #{@host}:#{@port}")
+        Ignis.logger.info("Model: #{@model_name} (#{@generator.model.num_parameters} params)")
+
+        @batch_processor.start!
+
+        if defined?(WNAIS::Server)
+          start_wnais_server!
+        else
+          start_tcp_server!
+        end
+      end
+
+      # Stop the server.
+      # @return [void]
+      def stop!
+        @batch_processor.stop!
+        @server_thread&.kill
+        Ignis.logger.info("Server stopped")
+      end
+
+      private
+
+      # Start using WNAIS HTTP/2 infrastructure.
+      def start_wnais_server!
+        Ignis.logger.info("Using WNAIS HTTP/2 transport")
+        # WNAIS integration point — register routes
+        if defined?(Ignis::Shared::EventBus)
+          Ignis::Shared::EventBus.subscribe(:http_request) do |request|
+            handle_request(request)
+          end
+        end
+      end
+
+      # Fallback TCP server for standalone use.
+      def start_tcp_server!
+        require "socket"
+        require "uri"
+
+        @server_thread = Thread.new do
+          server = TCPServer.new(@host, @port)
+          Ignis.logger.info("TCP server listening on #{@host}:#{@port}")
+
+          loop do
+            client = server.accept
+            Thread.new(client) do |sock|
+              begin
+                handle_tcp_client(sock)
+              rescue => e
+                Ignis.logger.error("Client error: #{e.message}")
+              ensure
+                sock.close
+              end
+            end
+          end
+        end
+      end
+
+      # Handle a TCP client connection.
+      # @param sock [TCPSocket]
+      def handle_tcp_client(sock)
+        request_line = sock.gets
+        return unless request_line
+
+        method, path, _version = request_line.strip.split(" ")
+
+        # Read headers
+        headers = {}
+        while (line = sock.gets) && line.strip != ""
+          key, value = line.strip.split(": ", 2)
+          headers[key.downcase] = value
+        end
+
+        # Read body
+        body = nil
+        if headers["content-length"]
+          body = sock.read(headers["content-length"].to_i)
+        end
+
+        # Route
+        response = route_request(method, path, body, headers)
+
+        # Send response
+        sock.print "HTTP/1.1 #{response[:status]}\r\n"
+        sock.print "Content-Type: application/json\r\n"
+        sock.print "Access-Control-Allow-Origin: *\r\n"
+        sock.print "Content-Length: #{response[:body].length}\r\n"
+        sock.print "\r\n"
+        sock.print response[:body]
+      end
+
+      # Route request to handler.
+      # @param method [String]
+      # @param path [String]
+      # @param body [String, nil]
+      # @param headers [Hash]
+      # @return [Hash] :status, :body
+      def route_request(method, path, body, headers)
+        case [method, path]
+        when ["POST", "/v1/completions"]
+          handle_completions(body)
+        when ["POST", "/v1/chat/completions"]
+          handle_chat_completions(body)
+        when ["GET", "/v1/models"]
+          handle_models
+        when ["POST", "/v1/embeddings"]
+          handle_embeddings(body)
+        else
+          { status: "404 Not Found", body: JSON.generate({ error: "Not found" }) }
+        end
+      end
+
+      # POST /v1/completions
+      def handle_completions(body)
+        params = JSON.parse(body)
+        prompt = params["prompt"] || ""
+        max_tokens = params["max_tokens"] || 128
+        temperature = params["temperature"] || 0.7
+        top_p = params["top_p"] || 0.9
+
+        text = @generator.generate(prompt,
+                                    max_tokens: max_tokens,
+                                    temperature: temperature,
+                                    top_p: top_p)
+
+        response = {
+          id: "cmpl-#{SecureRandom.hex(12)}",
+          object: "text_completion",
+          created: Time.now.to_i,
+          model: @model_name,
+          choices: [{
+            text: text[prompt.length..],
+            index: 0,
+            finish_reason: "stop"
+          }],
+          usage: {
+            prompt_tokens: @generator.tokenizer.encode(prompt).length,
+            completion_tokens: max_tokens,
+            total_tokens: @generator.tokenizer.encode(text).length
+          }
+        }
+
+        { status: "200 OK", body: JSON.generate(response) }
+      rescue => e
+        { status: "500 Internal Server Error", body: JSON.generate({ error: e.message }) }
+      end
+
+      # POST /v1/chat/completions
+      def handle_chat_completions(body)
+        params = JSON.parse(body)
+        messages = params["messages"] || []
+
+        # Convert chat messages to prompt string
+        prompt = messages.map { |m| "#{m['role']}: #{m['content']}" }.join("\n")
+        prompt += "\nassistant: "
+
+        max_tokens = params["max_tokens"] || 128
+        temperature = params["temperature"] || 0.7
+
+        text = @generator.generate(prompt,
+                                    max_tokens: max_tokens,
+                                    temperature: temperature)
+        response_text = text[prompt.length..]
+
+        response = {
+          id: "chatcmpl-#{SecureRandom.hex(12)}",
+          object: "chat.completion",
+          created: Time.now.to_i,
+          model: @model_name,
+          choices: [{
+            index: 0,
+            message: { role: "assistant", content: response_text },
+            finish_reason: "stop"
+          }],
+          usage: {
+            prompt_tokens: @generator.tokenizer.encode(prompt).length,
+            completion_tokens: @generator.tokenizer.encode(response_text).length,
+            total_tokens: @generator.tokenizer.encode(text).length
+          }
+        }
+
+        { status: "200 OK", body: JSON.generate(response) }
+      rescue => e
+        { status: "500 Internal Server Error", body: JSON.generate({ error: e.message }) }
+      end
+
+      # GET /v1/models
+      def handle_models
+        response = {
+          object: "list",
+          data: [{
+            id: @model_name,
+            object: "model",
+            created: Time.now.to_i,
+            owned_by: "nnw",
+            permission: [],
+            root: @model_name,
+            parent: nil
+          }]
+        }
+
+        { status: "200 OK", body: JSON.generate(response) }
+      end
+
+      # POST /v1/embeddings
+      def handle_embeddings(body)
+        params = JSON.parse(body)
+        input = params["input"]
+        inputs = input.is_a?(Array) ? input : [input]
+
+        embeddings = inputs.map.with_index do |text, i|
+          emb = @generator.embed(text)
+          { object: "embedding", embedding: emb, index: i }
+        end
+
+        response = {
+          object: "list",
+          data: embeddings,
+          model: @model_name,
+          usage: { prompt_tokens: inputs.sum { |t| @generator.tokenizer.encode(t).length }, total_tokens: 0 }
+        }
+
+        { status: "200 OK", body: JSON.generate(response) }
+      rescue => e
+        { status: "500 Internal Server Error", body: JSON.generate({ error: e.message }) }
+      end
+    end
+  end
+end

data/lib/nnw/ai/tokenizer.rb

@@ -0,0 +1,413 @@
+# frozen_string_literal: true
+
+require "json"
+require "set"
+
+module Ignis
+  module AI
+    # Hybrid tokenizer: tries native HuggingFace tokenizers DLL first,
+    # falls back to pure Ruby BPE. Best of both worlds.
+    #
+    # Native DLL: tokenizers_ruby.dll / libtokenizers.so / libtokenizers.dylib
+    # Search paths: model dir, Ignis lib root, system PATH, vcpkg
+    class Tokenizer
+      # @return [Boolean] whether using native backend
+      attr_reader :native_backend
+      # @return [Integer] vocabulary size
+      attr_reader :vocab_size
+
+      # @return [Hash{String => Integer}] token to id
+      attr_reader :token_to_id
+
+      # @return [Hash{Integer => String}] id to token
+      attr_reader :id_to_token
+
+      # @return [Hash{String => Integer}] special tokens
+      attr_reader :special_tokens
+
+      # @return [Set<Integer>] ids of special tokens (used by inference EOS checks)
+      attr_reader :special_token_ids
+
+      # @param config_path [String] path to tokenizer.json
+      def initialize(config_path)
+        raise ArgumentError, "tokenizer.json not found: #{config_path}" unless File.exist?(config_path)
+
+        @config_path = config_path
+        @native_backend = false
+        @native_handle = nil
+
+        # Try native DLL first
+        native_dll = find_native_dll(File.dirname(config_path))
+        if native_dll
+          begin
+            load_native_backend(native_dll, config_path)
+            @native_backend = true
+            Ignis.logger.info("Tokenizer: using native backend (#{File.basename(native_dll)})")
+          rescue => e
+            Ignis.logger.warn("Tokenizer: native backend failed (#{e.message}), falling back to Ruby BPE")
+            @native_backend = false
+          end
+        end
+
+        # Always load Ruby config for metadata (special tokens, vocab size)
+        # even when native backend is in use
+        config = JSON.parse(File.read(config_path))
+        load_from_config(config)
+      end
+
+      # Load from HuggingFace model directory.
+      # @param dir [String] directory containing tokenizer.json
+      # @return [Tokenizer]
+      def self.from_pretrained(dir)
+        path = File.join(dir, "tokenizer.json")
+        new(path)
+      end
+
+      # Encode text to token ids.
+      # Uses native backend if available, else pure Ruby BPE.
+      # @param text [String]
+      # @param add_special_tokens [Boolean]
+      # @return [Array<Integer>]
+      def encode(text, add_special_tokens: true)
+        if @native_backend && @native_encode
+          native_encode(text, add_special_tokens)
+        else
+          ruby_encode(text, add_special_tokens)
+        end
+      end
+
+      # Decode token ids to text.
+      # @param ids [Array<Integer>]
+      # @param skip_special_tokens [Boolean]
+      # @return [String]
+      def decode(ids, skip_special_tokens: true)
+        tokens = ids.filter_map do |id|
+          next nil if skip_special_tokens && @special_token_ids.include?(id)
+          @id_to_token[id]
+        end
+        text = tokens.join("")
+        @byte_level ? decode_byte_level(text) : text
+      end
+
+      # Batch encode.
+      # @param texts [Array<String>]
+      # @param add_special_tokens [Boolean]
+      # @return [Array<Array<Integer>>]
+      def encode_batch(texts, add_special_tokens: true)
+        texts.map { |t| encode(t, add_special_tokens: add_special_tokens) }
+      end
+
+      # Batch decode.
+      # @param id_sequences [Array<Array<Integer>>]
+      # @param skip_special_tokens [Boolean]
+      # @return [Array<String>]
+      def decode_batch(id_sequences, skip_special_tokens: true)
+        id_sequences.map { |ids| decode(ids, skip_special_tokens: skip_special_tokens) }
+      end
+
+      # Encode and return GPU Tensor (int32).
+      # @param text [String]
+      # @param device_id [Integer]
+      # @return [Tensor]
+      def encode_to_tensor(text, device_id: 0)
+        ids = encode(text)
+        nv = Ignis::Shared::NvArray.new(shape: [1, ids.length], dtype: :int32, device_id: device_id)
+        nv.from_host(ids)
+        Tensor.new(data: nv, requires_grad: false)
+      end
+
+      private
+
+      # Load from tokenizer.json.
+      # @param config [Hash]
+      def load_from_config(config)
+        model = config["model"] || {}
+        @token_to_id = {}
+        @id_to_token = {}
+        @merges = []
+        @merge_ranks = {}
+        @byte_level = false
+
+        # Vocab
+        (model["vocab"] || {}).each do |token, id|
+          @token_to_id[token] = id
+          @id_to_token[id] = token
+        end
+
+        # Merges
+        (model["merges"] || []).each_with_index do |merge_str, rank|
+          pair = merge_str.split(" ", 2)
+          @merges << pair
+          @merge_ranks[pair.join(" ")] = rank
+        end
+
+        @vocab_size = @token_to_id.size
+
+        # Special tokens
+        @special_tokens = {}
+        @special_token_ids = Set.new
+        (config["added_tokens"] || []).each do |info|
+          content = info["content"]
+          id = info["id"]
+          if info["special"]
+            @special_tokens[content] = id
+            @special_token_ids << id
+          end
+          @token_to_id[content] = id
+          @id_to_token[id] = content
+        end
+
+        # Detect byte-level pre-tokenizer
+        pre_tok = config["pre_tokenizer"]
+        if pre_tok
+          @byte_level = (pre_tok["type"] == "ByteLevel") ||
+                        (pre_tok.dig("pretokenizers", 0, "type") == "ByteLevel")
+        end
+
+        # Build byte-level lookup tables (GPT-2 style)
+        build_byte_encoder if @byte_level
+      end
+
+      # Ruby BPE encode (fallback when no native DLL).
+      # @param text [String]
+      # @param add_special_tokens [Boolean]
+      # @return [Array<Integer>]
+      def ruby_encode(text, add_special_tokens)
+        words = pre_tokenize(text)
+        ids = []
+        words.each do |word|
+          ids.concat(bpe_encode(word))
+        end
+        ids
+      end
+
+      # GPT-2 regex-style pre-tokenization: split on word boundaries.
+      # @param text [String]
+      # @return [Array<String>]
+      def pre_tokenize(text)
+        if @byte_level
+          pattern = /'s|'t|'re|'ve|'m|'ll|'d| ?\w+| ?\d+| ?[^\s\w\d]+|\s+/
+          tokens = text.scan(pattern)
+          tokens.map { |word| encode_byte_level(word) }
+        else
+          text.split(/(\s+)/).reject(&:empty?)
+        end
+      end
+
+      # BPE encode a single pre-tokenized word.
+      # @param word [String]
+      # @return [Array<Integer>]
+      def bpe_encode(word)
+        chars = word.chars
+
+        if chars.length == 1
+          id = @token_to_id[chars[0]]
+          return id ? [id] : encode_unknown(chars[0])
+        end
+
+        pairs = get_pairs(chars)
+        loop do
+          best_pair = nil
+          best_rank = Float::INFINITY
+
+          pairs.each do |pair|
+            key = pair.join(" ")
+            rank = @merge_ranks[key]
+            if rank && rank < best_rank
+              best_rank = rank
+              best_pair = pair
+            end
+          end
+
+          break unless best_pair
+
+          chars = apply_merge(chars, best_pair)
+          break if chars.length == 1
+
+          pairs = get_pairs(chars)
+        end
+
+        chars.map do |token|
+          id = @token_to_id[token]
+          id || encode_unknown(token)
+        end.flatten
+      end
+
+      # @param tokens [Array<String>]
+      # @return [Array<Array<String>>]
+      def get_pairs(tokens)
+        pairs = []
+        (0...tokens.length - 1).each do |i|
+          pairs << [tokens[i], tokens[i + 1]]
+        end
+        pairs.uniq
+      end
+
+      # @param tokens [Array<String>]
+      # @param pair [Array<String>]
+      # @return [Array<String>]
+      def apply_merge(tokens, pair)
+        merged = pair.join("")
+        result = []
+        i = 0
+        while i < tokens.length
+          if i < tokens.length - 1 && tokens[i] == pair[0] && tokens[i + 1] == pair[1]
+            result << merged
+            i += 2
+          else
+            result << tokens[i]
+            i += 1
+          end
+        end
+        result
+      end
+
+      # @param token [String]
+      # @return [Array<Integer>]
+      def encode_unknown(token)
+        token.bytes.filter_map { |b| @token_to_id[@byte_encoder[b]] if @byte_encoder }
+      end
+
+      def build_byte_encoder
+        @byte_encoder = {}
+        @byte_decoder = {}
+
+        bs = (33..126).to_a + (161..172).to_a + (174..255).to_a
+        cs = bs.dup
+        n = 0
+        (0..255).each do |b|
+          unless bs.include?(b)
+            bs << b
+            cs << 256 + n
+            n += 1
+          end
+        end
+
+        bs.zip(cs).each do |b, c|
+          @byte_encoder[b] = c.chr(Encoding::UTF_8)
+          @byte_decoder[c.chr(Encoding::UTF_8)] = b
+        end
+      end
+
+      # @param text [String]
+      # @return [String]
+      def encode_byte_level(text)
+        text.bytes.map { |b| @byte_encoder[b] }.join("")
+      end
+
+      # @param text [String]
+      # @return [String]
+      def decode_byte_level(text)
+        bytes = text.chars.map { |c| @byte_decoder[c] || c.ord }.pack("C*")
+        bytes.force_encoding(Encoding::UTF_8)
+      end
+
+      # -------------------------------------------------------------------
+      # Native DLL Backend (HuggingFace tokenizers compiled from Rust)
+      # -------------------------------------------------------------------
+
+      # Search for native tokenizers library.
+      # @param model_dir [String] directory containing the model
+      # @return [String, nil] path to DLL, or nil
+      def find_native_dll(model_dir)
+        candidates = native_dll_candidates(model_dir)
+        candidates.find { |path| File.exist?(path) }
+      end
+
+      # Generate list of candidate DLL paths by platform.
+      # @param model_dir [String]
+      # @return [Array<String>]
+      def native_dll_candidates(model_dir)
+        paths = []
+
+        if Ignis::Platform.windows?
+          dll_name = "tokenizers_ruby.dll"
+          paths << File.join(model_dir, dll_name)
+          paths << File.join(__dir__, "..", "..", "..", "ext", dll_name)
+          paths << File.join(ENV["USERPROFILE"] || "C:\\Users\\#{ENV['USERNAME']}", "vcpkg", "installed", "x64-windows", "bin", dll_name)
+          # Check PATH
+          (ENV["PATH"] || "").split(";").each do |dir|
+            paths << File.join(dir, dll_name)
+          end
+        else
+          # Linux / macOS
+          ext = Ignis::Platform.macos? ? "dylib" : "so"
+          lib_name = "libtokenizers.#{ext}"
+          paths << File.join(model_dir, lib_name)
+          paths << File.join(__dir__, "..", "..", "..", "ext", lib_name)
+          paths << "/usr/local/lib/#{lib_name}"
+          paths << "/usr/lib/#{lib_name}"
+          (ENV["LD_LIBRARY_PATH"] || "").split(":").each do |dir|
+            paths << File.join(dir, lib_name)
+          end
+        end
+
+        paths
+      end
+
+      # Load native backend via Fiddle.
+      # @param dll_path [String]
+      # @param config_path [String]
+      # @return [void]
+      def load_native_backend(dll_path, config_path)
+        require "fiddle"
+        require "fiddle/import"
+
+        @native_lib = Fiddle.dlopen(dll_path)
+
+        # Bind tokenizer_from_file(path) → handle
+        @native_from_file = Fiddle::Function.new(
+          @native_lib["tokenizer_from_file"],
+          [Fiddle::TYPE_VOIDP],
+          Fiddle::TYPE_VOIDP
+        )
+
+        # Bind tokenizer_encode(handle, text, len, add_special, out_ids, out_len) → int
+        @native_encode_fn = Fiddle::Function.new(
+          @native_lib["tokenizer_encode"],
+          [Fiddle::TYPE_VOIDP, Fiddle::TYPE_VOIDP, Fiddle::TYPE_INT,
+           Fiddle::TYPE_INT, Fiddle::TYPE_VOIDP, Fiddle::TYPE_VOIDP],
+          Fiddle::TYPE_INT
+        )
+
+        # Bind tokenizer_decode(handle, ids, len, out_text, out_len) → int
+        @native_decode_fn = Fiddle::Function.new(
+          @native_lib["tokenizer_decode"],
+          [Fiddle::TYPE_VOIDP, Fiddle::TYPE_VOIDP, Fiddle::TYPE_INT,
+           Fiddle::TYPE_VOIDP, Fiddle::TYPE_VOIDP],
+          Fiddle::TYPE_INT
+        )
+
+        # Load tokenizer
+        config_ptr = Fiddle::Pointer.to_ptr(config_path)
+        @native_handle = @native_from_file.call(config_ptr)
+        raise "Failed to load tokenizer from #{config_path}" if @native_handle.null?
+
+        @native_encode = true
+        @native_decode = true
+      end
+
+      # Encode via native backend.
+      # @param text [String]
+      # @param add_special_tokens [Boolean]
+      # @return [Array<Integer>]
+      def native_encode(text, add_special_tokens)
+        max_tokens = text.length * 4  # Upper bound
+        out_ids = Fiddle::Pointer.malloc(max_tokens * 4, Fiddle::RUBY_FREE)
+        out_len = Fiddle::Pointer.malloc(4, Fiddle::RUBY_FREE)
+
+        text_ptr = Fiddle::Pointer.to_ptr(text)
+        status = @native_encode_fn.call(@native_handle, text_ptr, text.bytesize,
+                                         add_special_tokens ? 1 : 0, out_ids, out_len)
+
+        if status == 0
+          n = out_len[0, 4].unpack1("l")
+          out_ids[0, n * 4].unpack("l*")
+        else
+          # Fallback to Ruby on error
+          ruby_encode(text, add_special_tokens)
+        end
+      end
+    end
+  end
+end