red-candle 1.0.0.pre.7 → 1.0.1

data/lib/candle/build_info.rb

@@ -7,6 +7,7 @@ module Candle
       return unless ENV['CANDLE_VERBOSE'] || ENV['CANDLE_DEBUG'] || $DEBUG
       
       if info["cuda_available"] == false
+        # :nocov:
         # Check if CUDA could be available on the system
         cuda_potentially_available = ENV['CUDA_ROOT'] || ENV['CUDA_PATH'] || 
                                    File.exist?('/usr/local/cuda') || File.exist?('/opt/cuda')
@@ -18,6 +19,7 @@ module Candle
           warn "  CANDLE_ENABLE_CUDA=1 gem install red-candle"
           warn "=" * 80
         end
+        # :nocov:
       end
     end
     

data/lib/candle/device_utils.rb

@@ -7,6 +7,7 @@ module Candle
         # Try Metal first (for Mac users)
         Device.metal
       rescue
+        # :nocov:
         begin
           # Try CUDA next (for NVIDIA GPU users)
           Device.cuda
@@ -14,6 +15,7 @@ module Candle
           # Fall back to CPU
           Device.cpu
         end
+        # :nocov:
       end
     end
   end

data/lib/candle/ner.rb

@@ -0,0 +1,345 @@
+# frozen_string_literal: true
+
+module Candle
+  # Named Entity Recognition (NER) for token classification
+  #
+  # This class provides methods to extract named entities from text using
+  # pre-trained BERT-based models. It supports standard NER labels like
+  # PER (person), ORG (organization), LOC (location), and can be extended
+  # with custom entity types.
+  #
+  # @example Load a pre-trained NER model
+  #   ner = Candle::NER.from_pretrained("Babelscape/wikineural-multilingual-ner")
+  #
+  # @example Load a model with a specific tokenizer
+  #   ner = Candle::NER.from_pretrained("dslim/bert-base-NER", tokenizer: "bert-base-cased")
+  #
+  # @example Extract entities from text
+  #   entities = ner.extract_entities("Apple Inc. was founded by Steve Jobs in Cupertino.")
+  #   # => [
+  #   #   { text: "Apple Inc.", label: "ORG", start: 0, end: 10, confidence: 0.99 },
+  #   #   { text: "Steve Jobs", label: "PER", start: 26, end: 36, confidence: 0.98 },
+  #   #   { text: "Cupertino", label: "LOC", start: 40, end: 49, confidence: 0.97 }
+  #   # ]
+  #
+  # @example Get token-level predictions
+  #   tokens = ner.predict_tokens("John works at Google")
+  #   # Returns detailed token-by-token predictions with confidence scores
+  class NER
+    class << self
+      # Load a pre-trained NER model from HuggingFace
+      #
+      # @param model_id [String] HuggingFace model ID (e.g., "dslim/bert-base-NER")
+      # @param device [Device, nil] Device to run on (defaults to best available)
+      # @param tokenizer [String, nil] Tokenizer model ID to use (defaults to same as model_id)
+      # @return [NER] NER instance
+      def from_pretrained(model_id, device: nil, tokenizer: nil)
+        new(model_id, device, tokenizer)
+      end
+      
+      # Popular pre-trained models for different domains
+      def suggested_models
+        {
+          general: {
+            model: "Babelscape/wikineural-multilingual-ner",
+            note: "Has tokenizer.json"
+          },
+          general_alt: {
+            model: "dslim/bert-base-NER",
+            tokenizer: "bert-base-cased",
+            note: "Requires separate tokenizer"
+          },
+          multilingual: {
+            model: "Davlan/bert-base-multilingual-cased-ner-hrl",
+            note: "Check tokenizer availability"
+          },
+          biomedical: {
+            model: "dmis-lab/biobert-base-cased-v1.2",
+            note: "May require specific tokenizer"
+          },
+          clinical: {
+            model: "emilyalsentzer/Bio_ClinicalBERT",
+            note: "May require specific tokenizer"
+          },
+          scientific: {
+            model: "allenai/scibert_scivocab_uncased",
+            note: "May require specific tokenizer"
+          }
+        }
+      end
+    end
+    
+    # Create an alias for the native method
+    alias_method :_extract_entities, :extract_entities
+    
+    # Extract entities from text
+    #
+    # @param text [String] The text to analyze
+    # @param confidence_threshold [Float] Minimum confidence score (default: 0.9)
+    # @return [Array<Hash>] Array of entity hashes with text, label, start, end, confidence
+    def extract_entities(text, confidence_threshold: 0.9)
+      # Call the native method with positional arguments
+      _extract_entities(text, confidence_threshold)
+    end
+    
+    # Get available entity types
+    #
+    # @return [Array<String>] List of entity types (without B-/I- prefixes)
+    def entity_types
+      return @entity_types if @entity_types
+      
+      label_config = labels
+      @entity_types = label_config["label2id"].keys
+        .reject { |l| l == "O" }
+        .map { |l| l.sub(/^[BI]-/, "") }
+        .uniq
+        .sort
+    end
+    
+    # Check if model supports a specific entity type
+    #
+    # @param entity_type [String] Entity type to check (e.g., "GENE", "PER")
+    # @return [Boolean] Whether the model recognizes this entity type
+    def supports_entity?(entity_type)
+      entity_types.include?(entity_type.upcase)
+    end
+    
+    # Extract entities of a specific type
+    #
+    # @param text [String] The text to analyze
+    # @param entity_type [String] Entity type to extract (e.g., "PER", "ORG")
+    # @param confidence_threshold [Float] Minimum confidence score
+    # @return [Array<Hash>] Filtered entities of the specified type
+    def extract_entity_type(text, entity_type, confidence_threshold: 0.9)
+      entities = extract_entities(text, confidence_threshold: confidence_threshold)
+      entities.select { |e| e["label"] == entity_type.upcase }
+    end
+    
+    # Analyze text and return both entities and token predictions
+    #
+    # @param text [String] The text to analyze
+    # @param confidence_threshold [Float] Minimum confidence for entities
+    # @return [Hash] Hash with :entities and :tokens keys
+    def analyze(text, confidence_threshold: 0.9)
+      {
+        entities: extract_entities(text, confidence_threshold: confidence_threshold),
+        tokens: predict_tokens(text)
+      }
+    end
+    
+    # Get a formatted string representation of entities
+    #
+    # @param text [String] The text to analyze
+    # @param confidence_threshold [Float] Minimum confidence score
+    # @return [String] Formatted output with entities highlighted
+    def format_entities(text, confidence_threshold: 0.9)
+      entities = extract_entities(text, confidence_threshold: confidence_threshold)
+      return text if entities.empty?
+      
+      # Sort by start position (reverse for easier insertion)
+      entities.sort_by! { |e| -e["start"] }
+      
+      result = text.dup
+      entities.each do |entity|
+        label = "[#{entity['label']}:#{entity['confidence'].round(2)}]"
+        result.insert(entity["end"], label)
+      end
+      
+      result
+    end
+    
+    # Get model information
+    #
+    # @return [String] Model description
+    def inspect
+      "#<Candle::NER #{model_info}>"
+    end
+    
+    alias to_s inspect
+  end
+  
+  # Pattern-based entity recognizer for custom entities
+  class PatternEntityRecognizer
+    attr_reader :patterns, :entity_type
+    
+    def initialize(entity_type, patterns = [])
+      @entity_type = entity_type
+      @patterns = patterns
+    end
+    
+    # Add a pattern (String or Regexp)
+    def add_pattern(pattern)
+      @patterns << pattern
+      self
+    end
+    
+    # Recognize entities using patterns
+    def recognize(text, tokenizer = nil)
+      entities = []
+      
+      @patterns.each do |pattern|
+        regex = pattern.is_a?(Regexp) ? pattern : Regexp.new(pattern)
+        
+        text.scan(regex) do |match|
+          match_text = $&
+          match_start = $~.offset(0)[0]
+          match_end = $~.offset(0)[1]
+          
+          entities << {
+            "text" => match_text,
+            "label" => @entity_type,
+            "start" => match_start,
+            "end" => match_end,
+            "confidence" => 1.0,
+            "source" => "pattern"
+          }
+        end
+      end
+      
+      entities
+    end
+  end
+  
+  # Gazetteer-based entity recognizer
+  class GazetteerEntityRecognizer
+    attr_reader :entity_type, :terms, :case_sensitive
+    
+    def initialize(entity_type, terms = [], case_sensitive: false)
+      @entity_type = entity_type
+      @case_sensitive = case_sensitive
+      @terms = build_term_set(terms)
+    end
+    
+    # Add terms to the gazetteer
+    def add_terms(terms)
+      terms = [terms] unless terms.is_a?(Array)
+      terms.each { |term| @terms.add(normalize_term(term)) }
+      self
+    end
+    
+    # Load terms from file
+    def load_from_file(filepath)
+      File.readlines(filepath).each do |line|
+        term = line.strip
+        add_terms(term) unless term.empty? || term.start_with?("#")
+      end
+      self
+    end
+    
+    # Recognize entities using the gazetteer
+    def recognize(text, tokenizer = nil)
+      entities = []
+      normalized_text = @case_sensitive ? text : text.downcase
+      
+      @terms.each do |term|
+        pattern = @case_sensitive ? term : term.downcase
+        pos = 0
+        
+        while (idx = normalized_text.index(pattern, pos))
+          # Check word boundaries
+          prev_char = idx > 0 ? text[idx - 1] : " "
+          next_char = idx + pattern.length < text.length ? text[idx + pattern.length] : " "
+          
+          if word_boundary?(prev_char) && word_boundary?(next_char)
+            entities << {
+              "text" => text[idx, pattern.length],
+              "label" => @entity_type,
+              "start" => idx,
+              "end" => idx + pattern.length,
+              "confidence" => 1.0,
+              "source" => "gazetteer"
+            }
+          end
+          
+          pos = idx + 1
+        end
+      end
+      
+      entities
+    end
+    
+    private
+    
+    def build_term_set(terms)
+      Set.new(terms.map { |term| normalize_term(term) })
+    end
+    
+    def normalize_term(term)
+      @case_sensitive ? term : term.downcase
+    end
+    
+    def word_boundary?(char)
+      char.match?(/\W/)
+    end
+  end
+  
+  # Hybrid NER that combines ML model with rules
+  class HybridNER
+    attr_reader :model_ner, :pattern_recognizers, :gazetteer_recognizers
+    
+    def initialize(model_id = nil, device: nil)
+      @model_ner = model_id ? NER.from_pretrained(model_id, device: device) : nil
+      @pattern_recognizers = []
+      @gazetteer_recognizers = []
+    end
+    
+    # Add a pattern-based recognizer
+    def add_pattern_recognizer(entity_type, patterns)
+      recognizer = PatternEntityRecognizer.new(entity_type, patterns)
+      @pattern_recognizers << recognizer
+      self
+    end
+    
+    # Add a gazetteer-based recognizer
+    def add_gazetteer_recognizer(entity_type, terms, **options)
+      recognizer = GazetteerEntityRecognizer.new(entity_type, terms, **options)
+      @gazetteer_recognizers << recognizer
+      self
+    end
+    
+    # Extract entities using all recognizers
+    def extract_entities(text, confidence_threshold: 0.9)
+      all_entities = []
+      
+      # Model-based entities
+      if @model_ner
+        model_entities = @model_ner.extract_entities(text, confidence_threshold: confidence_threshold)
+        all_entities.concat(model_entities)
+      end
+      
+      # Pattern-based entities
+      @pattern_recognizers.each do |recognizer|
+        pattern_entities = recognizer.recognize(text)
+        all_entities.concat(pattern_entities)
+      end
+      
+      # Gazetteer-based entities
+      @gazetteer_recognizers.each do |recognizer|
+        gazetteer_entities = recognizer.recognize(text)
+        all_entities.concat(gazetteer_entities)
+      end
+      
+      # Merge overlapping entities (prefer highest confidence)
+      merge_entities(all_entities)
+    end
+    
+    private
+    
+    def merge_entities(entities)
+      # Sort by start position and confidence (descending)
+      sorted = entities.sort_by { |e| [e["start"], -e["confidence"]] }
+      
+      merged = []
+      sorted.each do |entity|
+        # Check if entity overlaps with any already merged
+        overlaps = merged.any? do |existing|
+          entity["start"] < existing["end"] && entity["end"] > existing["start"]
+        end
+        
+        merged << entity unless overlaps
+      end
+      
+      merged.sort_by { |e| e["start"] }
+    end
+  end
+end

data/lib/candle/reranker.rb

@@ -10,7 +10,7 @@ module Candle
       _create(model_path, device)
     end
 
-    # Returns the embedding for a string using the specified pooling method.
+    # Returns documents ranked by relevance using the specified pooling method.
     # @param query [String] The input text
     # @param documents [Array<String>] The list of documents to compare against
     # @param pooling_method [String] Pooling method: "pooler", "cls", or "mean". Default: "pooler"

data/lib/candle/tensor.rb

@@ -14,6 +14,7 @@ module Candle
           begin
             values_f32.each { |value| yield value }
           rescue NoMethodError
+            # :nocov:
             # If values_f32 isn't available yet (not recompiled), fall back
             if device.to_s != "cpu"
               # Move to CPU to avoid Metal F32->F64 conversion issue
@@ -21,6 +22,7 @@ module Candle
             else
               values.each { |value| yield value }
             end
+            # :nocov:
           end
         else
           # For non-F32 dtypes, use regular values

data/lib/candle/tokenizer.rb

@@ -0,0 +1,139 @@
+# frozen_string_literal: true
+
+module Candle
+  # Tokenizer class for text tokenization
+  #
+  # This class provides methods to encode text into tokens and decode tokens back to text.
+  # It supports both single text and batch processing, with options for special tokens,
+  # padding, and truncation.
+  #
+  # @example Create a tokenizer from a pretrained model
+  #   tokenizer = Candle::Tokenizer.from_pretrained("bert-base-uncased")
+  #
+  # @example Encode and decode text
+  #   tokens = tokenizer.encode("Hello, world!")
+  #   text = tokenizer.decode(tokens)
+  #
+  # @example Batch encoding
+  #   texts = ["Hello", "World", "Test"]
+  #   batch_tokens = tokenizer.encode_batch(texts)
+  #
+  # @example Configure padding and truncation
+  #   padded_tokenizer = tokenizer.with_padding(length: 128)
+  #   truncated_tokenizer = tokenizer.with_truncation(512)
+  class Tokenizer
+    # These methods are implemented in Rust
+    # - from_file(path) - Load tokenizer from a JSON file
+    # - from_pretrained(model_id) - Load tokenizer from HuggingFace
+    # - encode(text, add_special_tokens = true) - Encode text to token IDs
+    # - encode_to_tokens(text, add_special_tokens = true) - Encode text to token strings
+    # - encode_with_tokens(text, add_special_tokens = true) - Get both IDs and tokens
+    # - encode_batch(texts, add_special_tokens = true) - Encode multiple texts to IDs
+    # - encode_batch_to_tokens(texts, add_special_tokens = true) - Encode multiple texts to tokens
+    # - decode(token_ids, skip_special_tokens = true) - Decode token IDs to text
+    # - id_to_token(token_id) - Get token string for a token ID
+    # - get_vocab(with_added_tokens = true) - Get vocabulary as hash
+    # - vocab_size(with_added_tokens = true) - Get vocabulary size
+    # - with_padding(options) - Create tokenizer with padding enabled
+    # - with_truncation(max_length) - Create tokenizer with truncation enabled
+    # - get_special_tokens - Get special tokens info
+    # - inspect - String representation
+    # - to_s - String representation
+
+    # The native methods accept positional arguments, but we provide keyword argument interfaces
+    # for better Ruby ergonomics. We need to call the native methods with positional args.
+    
+    alias_method :_native_encode, :encode
+    alias_method :_native_encode_to_tokens, :encode_to_tokens
+    alias_method :_native_encode_with_tokens, :encode_with_tokens
+    alias_method :_native_encode_batch, :encode_batch
+    alias_method :_native_encode_batch_to_tokens, :encode_batch_to_tokens
+    alias_method :_native_decode, :decode
+    alias_method :_native_get_vocab, :get_vocab
+    alias_method :_native_vocab_size, :vocab_size
+    alias_method :_native_with_padding, :with_padding
+
+    # Encode text with convenient keyword arguments
+    #
+    # @param text [String] The text to encode
+    # @param add_special_tokens [Boolean] Whether to add special tokens (default: true)
+    # @return [Array<Integer>] Token IDs
+    def encode(text, add_special_tokens: true)
+      _native_encode(text, add_special_tokens)
+    end
+
+    # Encode text into token strings (words/subwords)
+    #
+    # @param text [String] The text to encode
+    # @param add_special_tokens [Boolean] Whether to add special tokens (default: true)
+    # @return [Array<String>] Token strings
+    def encode_to_tokens(text, add_special_tokens: true)
+      _native_encode_to_tokens(text, add_special_tokens)
+    end
+
+    # Encode text and return both IDs and token strings
+    #
+    # @param text [String] The text to encode
+    # @param add_special_tokens [Boolean] Whether to add special tokens (default: true)
+    # @return [Hash] Hash with :ids and :tokens arrays
+    def encode_with_tokens(text, add_special_tokens: true)
+      _native_encode_with_tokens(text, add_special_tokens)
+    end
+
+    # Encode multiple texts with convenient keyword arguments
+    #
+    # @param texts [Array<String>] The texts to encode
+    # @param add_special_tokens [Boolean] Whether to add special tokens (default: true)
+    # @return [Array<Array<Integer>>] Arrays of token IDs
+    def encode_batch(texts, add_special_tokens: true)
+      _native_encode_batch(texts, add_special_tokens)
+    end
+
+    # Encode multiple texts into token strings
+    #
+    # @param texts [Array<String>] The texts to encode
+    # @param add_special_tokens [Boolean] Whether to add special tokens (default: true)
+    # @return [Array<Array<String>>] Arrays of token strings
+    def encode_batch_to_tokens(texts, add_special_tokens: true)
+      _native_encode_batch_to_tokens(texts, add_special_tokens)
+    end
+
+    # Decode token IDs with convenient keyword arguments
+    #
+    # @param token_ids [Array<Integer>] The token IDs to decode
+    # @param skip_special_tokens [Boolean] Whether to skip special tokens (default: true)
+    # @return [String] Decoded text
+    def decode(token_ids, skip_special_tokens: true)
+      _native_decode(token_ids, skip_special_tokens)
+    end
+
+    # Get vocabulary with convenient keyword arguments
+    #
+    # @param with_added_tokens [Boolean] Include added tokens (default: true)
+    # @return [Hash<String, Integer>] Token to ID mapping
+    def get_vocab(with_added_tokens: true)
+      _native_get_vocab(with_added_tokens)
+    end
+
+    # Get vocabulary size with convenient keyword arguments
+    #
+    # @param with_added_tokens [Boolean] Include added tokens (default: true)
+    # @return [Integer] Vocabulary size
+    def vocab_size(with_added_tokens: true)
+      _native_vocab_size(with_added_tokens)
+    end
+
+    # Create a new tokenizer with padding configuration
+    #
+    # @param options [Hash] Padding options
+    # @option options [Integer] :length Fixed length padding
+    # @option options [Boolean] :max_length Use batch longest padding
+    # @option options [String] :direction Padding direction ("left" or "right")
+    # @option options [Integer] :pad_id Padding token ID
+    # @option options [String] :pad_token Padding token string
+    # @return [Tokenizer] New tokenizer instance with padding enabled
+    def with_padding(**options)
+      _native_with_padding(options)
+    end
+  end
+end

data/lib/candle/version.rb

@@ -1,3 +1,5 @@
+# :nocov:
 module Candle
-  VERSION = "1.0.0.pre.7"
-end
+  VERSION = "1.0.1"
+end
+# :nocov:

data/lib/candle.rb

@@ -5,4 +5,6 @@ require_relative "candle/embedding_model_type"
 require_relative "candle/embedding_model"
 require_relative "candle/reranker"
 require_relative "candle/llm"
+require_relative "candle/tokenizer"
+require_relative "candle/ner"
 require_relative "candle/build_info"