red-candle 1.8.0.pre3-aarch64-linux

data/lib/candle/llm.rb

@@ -0,0 +1,595 @@
+require 'json'
+
+module Candle
+  class LLM
+    # Cache for EOS token to avoid repeated calls
+    def cached_eos_token
+      @cached_eos_token ||= begin
+        if respond_to?(:eos_token)
+          eos_token rescue nil
+        end
+      end
+    end
+    
+    # Get model-specific EOS tokens
+    def model_eos_tokens
+      @model_eos_tokens ||= begin
+        tokens = []
+        if model_eos = cached_eos_token
+          tokens << model_eos
+          # For Gemma, also include end_of_turn for chat scenarios and </s>
+          # Even though </s> is technically an HTML tag in Gemma's vocabulary,
+          # it seems to use it as a generation boundary in practice
+          if model_name.downcase.include?("gemma")
+            tokens << "<end_of_turn>"
+            tokens << "</s>"
+          end
+        else
+          # Fallback to common tokens only if model doesn't provide one
+          tokens = ["</s>", "<|endoftext|>", "<|im_end|>", "<end>"]
+        end
+        tokens.uniq
+      end
+    end
+    # Create a structured constraint from a JSON schema
+    # Uses the model's vocabulary with proper byte encoding handling
+    def constraint_from_schema(schema)
+      schema_str = schema.is_a?(String) ? schema : JSON.generate(schema)
+
+      # Extract the tokenizer source model ID for proper vocabulary loading
+      tokenizer_model = tokenizer_source_model
+      if tokenizer_model
+        begin
+          StructuredConstraint.from_schema_with_model(schema_str, tokenizer_model)
+        rescue RuntimeError => e
+          # Fall back to legacy method if from_pretrained fails
+          # (e.g., tokenizer doesn't have EOS token in expected format)
+          if e.message.include?("UnsupportedTokenizer")
+            StructuredConstraint.from_schema(schema_str, tokenizer)
+          else
+            raise
+          end
+        end
+      else
+        # Fall back to legacy method if we can't determine the model
+        StructuredConstraint.from_schema(schema_str, tokenizer)
+      end
+    end
+
+    # Create a structured constraint from a regex pattern
+    # Uses the model's vocabulary with proper byte encoding handling
+    def constraint_from_regex(pattern)
+      pattern_str = pattern.is_a?(Regexp) ? pattern.source : pattern.to_s
+
+      # Extract the tokenizer source model ID for proper vocabulary loading
+      tokenizer_model = tokenizer_source_model
+      if tokenizer_model
+        begin
+          StructuredConstraint.from_regex_with_model(pattern_str, tokenizer_model)
+        rescue RuntimeError => e
+          # Fall back to legacy method if from_pretrained fails
+          if e.message.include?("UnsupportedTokenizer")
+            StructuredConstraint.from_regex(pattern_str, tokenizer)
+          else
+            raise
+          end
+        end
+      else
+        # Fall back to legacy method if we can't determine the model
+        StructuredConstraint.from_regex(pattern_str, tokenizer)
+      end
+    end
+
+    private
+
+    # Get the model ID to use for vocabulary loading
+    # This handles GGUF models by extracting the tokenizer source
+    def tokenizer_source_model
+      opts = options rescue {}
+
+      # For GGUF models, use the tokenizer source if available
+      if opts["tokenizer_source"]
+        return opts["tokenizer_source"]
+      end
+
+      # For regular models, use the base model ID
+      if opts["base_model"]
+        return opts["base_model"]
+      end
+
+      # Try model_id but strip GGUF parts
+      model = opts["model_id"] || (model_id rescue nil)
+      return nil unless model
+
+      # Remove GGUF file suffix if present
+      if model.include?("@")
+        model = model.split("@").first
+      end
+
+      # For GGUF repos, try to guess the tokenizer source
+      if model.downcase.include?("gguf")
+        guessed = self.class.guess_tokenizer(model)
+        return guessed if guessed && guessed != model
+      end
+
+      model
+    end
+
+    public
+    
+    # Generate with regex constraint
+    def generate_regex(prompt, pattern:, stop_on_match: true, **options)
+      constraint = constraint_from_regex(pattern)
+      
+      # Configure generation with early stopping by default
+      config_opts = options.merge(
+        constraint: constraint,
+        stop_on_constraint_satisfaction: options.fetch(:stop_on_constraint_satisfaction, stop_on_match),
+        stop_on_match: stop_on_match
+      )
+      config = options[:config] || GenerationConfig.balanced(**config_opts)
+      
+      generate(prompt, config: config, reset_cache: options.fetch(:reset_cache, true))
+    end
+    
+    # Generate and parse structured output from a JSON schema
+    def generate_structured(prompt, schema:, **options)
+      constraint = constraint_from_schema(schema)
+      
+      # Configure generation with early stopping by default
+      config_opts = options.merge(
+        constraint: constraint,
+        stop_on_constraint_satisfaction: options.fetch(:stop_on_constraint_satisfaction, true)
+      )
+      config = options[:config] || GenerationConfig.balanced(**config_opts)
+      
+      result = generate(prompt, config: config, reset_cache: options.fetch(:reset_cache, true))
+      
+      # Try to parse as JSON
+      begin
+        # First, try to extract JSON if there's content after stop tokens
+        json_content = extract_json_content(result)
+        JSON.parse(json_content)
+      rescue JSON::ParserError => e
+        # Return the raw string if parsing fails
+        Candle.logger.warn "Generated output is not valid JSON: #{e.message}" if options[:warn_on_parse_error]
+        result
+      end
+    end
+    # Tokenizer registry for automatic detection
+    TOKENIZER_REGISTRY = {
+      # Exact model matches
+      "TheBloke/Mistral-7B-Instruct-v0.2-GGUF" => "mistralai/Mistral-7B-Instruct-v0.2",
+      "TheBloke/Mistral-7B-v0.1-GGUF" => "mistralai/Mistral-7B-v0.1",
+      "TheBloke/Llama-2-7B-Chat-GGUF" => "meta-llama/Llama-2-7b-chat-hf",
+      "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF" => "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+      
+      # Qwen official GGUF models
+      "Qwen/Qwen3-8B-GGUF" => "Qwen/Qwen3-8B",
+      "Qwen/Qwen3-4B-GGUF" => "Qwen/Qwen3-4B",
+      "Qwen/Qwen3-14B-GGUF" => "Qwen/Qwen3-14B",
+      "Qwen/Qwen3-32B-GGUF" => "Qwen/Qwen3-32B",
+      "Qwen/Qwen3-72B-GGUF" => "Qwen/Qwen3-72B",
+      
+      # Phi GGUF models
+      "TheBloke/phi-2-GGUF" => "microsoft/phi-2",
+      "microsoft/phi-4-gguf" => "microsoft/phi-4",
+      "bartowski/Phi-3.5-mini-instruct-GGUF" => "microsoft/Phi-3.5-mini-instruct",
+      
+      # Pattern-based fallbacks (evaluated in order)
+      :patterns => [
+        # Mistral models
+        [/mistral.*?7b.*?instruct.*?v0\.2/i, "mistralai/Mistral-7B-Instruct-v0.2"],
+        [/mistral.*?7b.*?instruct.*?v0\.1/i, "mistralai/Mistral-7B-Instruct-v0.1"],
+        [/mistral.*?7b/i, "mistralai/Mistral-7B-v0.1"],
+        
+        # Llama models
+        [/llama.*?3.*?8b/i, "meta-llama/Meta-Llama-3-8B"],
+        [/llama.*?3.*?70b/i, "meta-llama/Meta-Llama-3-70B"],
+        [/llama.*?2.*?7b.*?chat/i, "meta-llama/Llama-2-7b-chat-hf"],
+        [/llama.*?2.*?13b.*?chat/i, "meta-llama/Llama-2-13b-chat-hf"],
+        [/llama.*?2.*?70b.*?chat/i, "meta-llama/Llama-2-70b-chat-hf"],
+        [/tinyllama/i, "TinyLlama/TinyLlama-1.1B-Chat-v1.0"],
+        
+        # Gemma models
+        [/gemma.*?2.*?9b/i, "google/gemma-2-9b"],
+        [/gemma.*?2.*?2b/i, "google/gemma-2-2b"],
+        [/gemma.*?7b/i, "google/gemma-7b"],
+        [/gemma.*?2b/i, "google/gemma-2b"],
+        
+        # Qwen models
+        [/qwen.*?3.*?72b/i, "Qwen/Qwen3-72B"],
+        [/qwen.*?3.*?32b/i, "Qwen/Qwen3-32B"],
+        [/qwen.*?3.*?14b/i, "Qwen/Qwen3-14B"],
+        [/qwen.*?3.*?8b/i, "Qwen/Qwen3-8B"],
+        [/qwen.*?3.*?4b/i, "Qwen/Qwen3-4B"],
+        [/qwen.*?3.*?1\.8b/i, "Qwen/Qwen3-1.8B"],
+        [/qwen.*?3.*?0\.5b/i, "Qwen/Qwen3-0.5B"],
+        [/qwen.*?2\.5/i, "Qwen/Qwen2.5-0.5B"],
+        [/qwen.*?2/i, "Qwen/Qwen2-1.5B"],
+        [/qwen/i, "Qwen/Qwen-1_8B"],
+        
+        # Phi models (order matters - more specific patterns first)
+        [/phi.*?3\.5.*?mini/i, "microsoft/Phi-3.5-mini-instruct"],
+        [/phi.*?3.*?mini.*?4k/i, "microsoft/Phi-3-mini-4k-instruct"],
+        [/phi.*?3.*?medium/i, "microsoft/Phi-3-medium-4k-instruct"],
+        [/phi.*?3.*?small/i, "microsoft/Phi-3-small-8k-instruct"],
+        [/phi.*?3.*?mini/i, "microsoft/Phi-3-mini-4k-instruct"],
+        [/phi.*?3/i, "microsoft/Phi-3-mini-4k-instruct"],
+        [/phi-4/i, "microsoft/phi-4"],
+        [/phi.*?2/i, "microsoft/phi-2"],
+        [/phi.*?1\.5/i, "microsoft/phi-1_5"],
+        [/phi/i, "microsoft/phi-2"]
+      ]
+    }
+    
+    # Allow users to register custom tokenizer mappings
+    def self.register_tokenizer(model_pattern, tokenizer_id)
+      if model_pattern.is_a?(String)
+        TOKENIZER_REGISTRY[model_pattern] = tokenizer_id
+      elsif model_pattern.is_a?(Regexp)
+        TOKENIZER_REGISTRY[:patterns] ||= []
+        TOKENIZER_REGISTRY[:patterns].unshift([model_pattern, tokenizer_id])
+      else
+        raise ArgumentError, "model_pattern must be a String or Regexp"
+      end
+    end
+    
+    # Guess the tokenizer for a model
+    def self.guess_tokenizer(model_id)
+      # Check exact matches first
+      return TOKENIZER_REGISTRY[model_id] if TOKENIZER_REGISTRY[model_id]
+      
+      # Check patterns
+      if patterns = TOKENIZER_REGISTRY[:patterns]
+        patterns.each do |pattern, tokenizer|
+          return tokenizer if model_id.match?(pattern)
+        end
+      end
+      
+      # Default: try removing common GGUF suffixes
+      base_model = model_id.gsub(/-gguf|-q\d+_\w+$/i, "")
+      base_model
+    end
+    
+    # Chat interface — always returns a String
+    def chat(messages, **options)
+      prompt = apply_chat_template(messages)
+      generate(prompt, **options)
+    end
+
+    # Streaming chat interface
+    def chat_stream(messages, **options, &block)
+      prompt = apply_chat_template(messages)
+      generate_stream(prompt, **options, &block)
+    end
+
+    # Chat with tool calling — always returns a ToolCallResult
+    # Set execute: true to automatically run the tools (default: false)
+    def chat_with_tools(messages, tools:, execute: false, **options)
+      tool_prompt = build_tool_system_prompt(tools)
+      augmented = inject_tool_instructions(messages, tool_prompt)
+
+      raw_response = chat(augmented, **options)
+
+      result = ToolCallParser.parse(raw_response, available_tools: tools)
+
+      if result.has_tool_calls? && execute
+        tool_results = result.tool_calls.map do |tool_call|
+          tool = tools.find { |t| t.name == tool_call.name }
+          unless tool
+            next { tool_call: tool_call, result: nil, error: "Unknown tool: #{tool_call.name}" }
+          end
+
+          begin
+            output = tool.call(tool_call.arguments)
+            { tool_call: tool_call, result: output, error: nil }
+          rescue Exception => e
+            { tool_call: tool_call, result: nil, error: e.message }
+          end
+        end
+
+        ToolCallResult.new(
+          tool_calls: result.tool_calls,
+          tool_results: tool_results,
+          text_response: result.text_response,
+          raw_response: raw_response
+        )
+      else
+        ToolCallResult.new(
+          tool_calls: result.tool_calls,
+          tool_results: [],
+          text_response: result.has_tool_calls? ? result.text_response : raw_response,
+          raw_response: raw_response
+        )
+      end
+    end
+
+    # Inspect method for debugging and exploration
+    def inspect
+      opts = options rescue {}
+      
+      # Extract key information
+      model_type = opts["model_type"] || "Unknown"
+      device = opts["device"] || self.device.to_s rescue "unknown"
+      
+      # Build the inspect string
+      parts = ["#<Candle::LLM"]
+      
+      # Add base model or model_id
+      if opts["base_model"]
+        parts << "model=#{opts["base_model"]}"
+      elsif opts["model_id"]
+        parts << "model=#{opts["model_id"]}"
+      elsif respond_to?(:model_id)
+        parts << "model=#{model_id}"
+      end
+      
+      # Add GGUF file if present
+      if opts["gguf_file"]
+        parts << "gguf=#{opts["gguf_file"]}"
+      end
+      
+      # Add device
+      parts << "device=#{device}"
+      
+      # Add model type
+      parts << "type=#{model_type}"
+      
+      # Add architecture for GGUF models
+      if opts["architecture"]
+        parts << "arch=#{opts["architecture"]}"
+      end
+      
+      parts.join(" ") + ">"
+    end
+
+    def generate(prompt, config: GenerationConfig.balanced, reset_cache: true)
+      begin
+        _generate(prompt, config)
+      ensure
+        clear_cache if reset_cache
+      end
+    end
+
+    def generate_stream(prompt, config: GenerationConfig.balanced, reset_cache: true, &block)
+      begin
+        _generate_stream(prompt, config, &block)
+      ensure
+        clear_cache if reset_cache
+      end
+    end
+
+    def self.from_pretrained(model_id, device: Candle::Device.best, gguf_file: nil, tokenizer: nil)
+      model_str = if gguf_file
+        "#{model_id}@#{gguf_file}"
+      else
+        model_id
+      end
+      
+      # Handle GGUF models that need tokenizer
+      if model_str.downcase.include?("gguf") && tokenizer.nil?
+        # Try to load without tokenizer first
+        begin
+          _from_pretrained(model_str, device)
+        rescue => e
+          if e.message.include?("No tokenizer found")
+            # Auto-detect tokenizer
+            detected_tokenizer = guess_tokenizer(model_id)
+            Candle.logger.info "No tokenizer found in GGUF repo. Using tokenizer from: #{detected_tokenizer}"
+            model_str = "#{model_str}@@#{detected_tokenizer}"
+            _from_pretrained(model_str, device)
+          else
+            raise e
+          end
+        end
+      elsif tokenizer
+        # User specified tokenizer
+        model_str = "#{model_str}@@#{tokenizer}"
+        _from_pretrained(model_str, device)
+      else
+        # Non-GGUF model or GGUF with embedded tokenizer
+        _from_pretrained(model_str, device)
+      end
+    end
+
+    private
+
+    def build_tool_system_prompt(tools)
+      tool_defs = tools.map { |t| JSON.generate(t.to_tool_definition) }.join("\n\n")
+      "You are a helpful assistant with access to the following tools:\n\n" \
+        "#{tool_defs}\n\n" \
+        "When you need to use a tool, respond with a tool call in the following format:\n" \
+        "<tool_call>\n" \
+        "{\"name\": \"tool_name\", \"arguments\": {\"arg1\": \"value1\"}}\n" \
+        "</tool_call>\n\n" \
+        "If you don't need to use a tool, respond normally with text."
+    end
+
+    def inject_tool_instructions(messages, tool_prompt)
+      msgs = messages.map { |m| m.dup }
+      if msgs.first && msgs.first[:role] == "system"
+        msgs.first[:content] = "#{tool_prompt}\n\n#{msgs.first[:content]}"
+      else
+        msgs.unshift({ role: "system", content: tool_prompt })
+      end
+      msgs
+    end
+
+    # Extract JSON content from generated text, handling stop tokens and extra content
+    def extract_json_content(text)
+      # Remove any content after common stop tokens
+      cleaned = text
+      
+      # Check for EOS tokens and truncate at the first one found
+      model_eos_tokens.each do |token|
+        if idx = cleaned.index(token)
+          cleaned = cleaned[0...idx]
+        end
+      end
+      
+      # Try to find valid JSON boundaries
+      # First try a simple approach - find the first { or [ and match to its closing } or ]
+      start_idx = cleaned.index(/[\{\[]/)
+      return cleaned.strip unless start_idx
+      
+      # Extract from the start position
+      json_candidate = cleaned[start_idx..-1]
+      
+      # Try to find a valid JSON object or array
+      # This regex handles nested structures better
+      if json_candidate[0] == '{'
+        # Match a JSON object
+        bracket_count = 0
+        in_string = false
+        escape_next = false
+        
+        json_candidate.chars.each_with_index do |char, idx|
+          if !in_string
+            case char
+            when '{'
+              bracket_count += 1
+            when '}'
+              bracket_count -= 1
+              if bracket_count == 0
+                return json_candidate[0..idx]
+              end
+            when '"'
+              in_string = true unless escape_next
+            end
+          else
+            if char == '"' && !escape_next
+              in_string = false
+            end
+          end
+          
+          escape_next = (!escape_next && char == '\\')
+        end
+      elsif json_candidate[0] == '['
+        # Match a JSON array (similar logic)
+        bracket_count = 0
+        in_string = false
+        escape_next = false
+        
+        json_candidate.chars.each_with_index do |char, idx|
+          if !in_string
+            case char
+            when '['
+              bracket_count += 1
+            when ']'
+              bracket_count -= 1
+              if bracket_count == 0
+                return json_candidate[0..idx]
+              end
+            when '"'
+              in_string = true unless escape_next
+            end
+          else
+            if char == '"' && !escape_next
+              in_string = false
+            end
+          end
+          
+          escape_next = (!escape_next && char == '\\')
+        end
+      end
+      
+      # If no valid JSON structure found, return the cleaned string
+      cleaned.strip
+    end
+
+    # Legacy format messages method - kept for backward compatibility
+    # Use apply_chat_template for proper model-specific formatting
+    def format_messages(messages)
+      formatted = messages.map do |msg|
+        case msg[:role]
+        when "system"
+          "System: #{msg[:content]}"
+        when "user"
+          "User: #{msg[:content]}"
+        when "assistant"
+          "Assistant: #{msg[:content]}"
+        else
+          msg[:content]
+        end
+      end.join("\n\n")
+      
+      # Add a prompt for the assistant to respond
+      formatted + "\n\nAssistant:"
+    end
+  end
+
+  class GenerationConfig
+    # Convenience method to create config with overrides
+    def with(**overrides)
+      current_config = {
+        max_length: max_length,
+        temperature: temperature,
+        top_p: top_p,
+        top_k: top_k,
+        repetition_penalty: repetition_penalty,
+        seed: seed,
+        stop_sequences: stop_sequences,
+        include_prompt: include_prompt,
+        constraint: defined?(@constraint) ? @constraint : nil
+      }.compact
+      
+      self.class.new(current_config.merge(overrides))
+    end
+
+    # Create a deterministic configuration (temperature = 0, fixed seed)
+    def self.deterministic(**opts)
+      defaults = {
+        temperature: 0.0,
+        top_p: nil,
+        top_k: 1,
+        seed: 42
+      }
+      new(defaults.merge(opts))
+    end
+
+    # Create a creative configuration (higher temperature, random seed)
+    def self.creative(**opts)
+      defaults = {
+        temperature: 1.0,
+        top_p: 0.95,
+        top_k: 50,
+        repetition_penalty: 1.2
+      }
+      new(defaults.merge(opts))
+    end
+
+    # Create a balanced configuration (moderate temperature, random seed)
+    def self.balanced(**opts)
+      defaults = {
+        temperature: 0.7,
+        top_p: 0.9,
+        top_k: 40
+      }
+      new(defaults.merge(opts))
+    end
+    
+    # Inspect method for debugging and exploration
+    def inspect
+      opts = options rescue {}
+      
+      parts = ["#<Candle::GenerationConfig"]
+      
+      # Add key configuration parameters
+      parts << "temp=#{opts["temperature"]}" if opts["temperature"]
+      parts << "max=#{opts["max_length"]}" if opts["max_length"]
+      parts << "top_p=#{opts["top_p"]}" if opts["top_p"]
+      parts << "top_k=#{opts["top_k"]}" if opts["top_k"]
+      parts << "seed=#{opts["seed"]}" if opts["seed"]
+      
+      # Add flags
+      flags = []
+      flags << "debug" if opts["debug_tokens"]
+      flags << "constraint" if opts["has_constraint"]
+      flags << "stop_on_match" if opts["stop_on_match"]
+      parts << "flags=[#{flags.join(",")}]" if flags.any?
+      
+      parts.join(" ") + ">"
+    end
+  end
+end