Run llama bench in all models from a directory and reports as a CSV

local_bench.rb

#!/usr/bin/env ruby
require 'json'
require 'csv'

def find_gguf_files(dir)
  return [] unless Dir.exist?(dir)

  Dir.glob(File.join(dir, '**', '*.gguf'))
     .reject { |f| f.include?('mmproj') }
end

def parse_model_name(filename)
  basename = File.basename(filename, '.gguf')
  basename.gsub(/[_-]/, ' ').strip
end

def get_file_size(path)
  size = File.size(path)
  if size >= 1_073_741_824
    format('%.2f GiB', size / 1_073_741_824.0)
  else
    format('%.2f MiB', size / 1_048_576.0)
  end
end

def ask_user(prompt)
  print prompt
  $stdin.gets.chomp.downcase
end

def run_bench(gguf_path, test_configs)
  results = []
  model_name = parse_model_name(gguf_path)
  file_size = get_file_size(gguf_path)

  test_configs.each do |config|
    test_name = config[:name]
    args = config[:args]

    cmd = "llama-bench -m \"#{gguf_path}\" #{args} -o json 2>/dev/null"
    output = `#{cmd}`

    begin
      json_data = JSON.parse(output)
      if json_data.is_a?(Array) && !json_data.empty?
        entry = json_data.first
        tps = entry['avg_ts'] || entry['t/s'] || 0
        results << {
          model: model_name,
          size: file_size,
          test: test_name,
          t_s: format('%.2f ± %.2f', tps, entry['std_ts'] || 0)
        }
      end
    rescue JSON::ParserError
      warn "Failed to parse output for #{model_name} (#{test_name})"
    end
  end

  results
end

if ARGV.empty?
  warn 'Usage: ruby bench_local.rb <models-directory>'
  exit 1
end

models_dir = ARGV[0]
unless Dir.exist?(models_dir)
  warn "Error: directory #{models_dir} does not exist"
  exit 1
end

gguf_files = find_gguf_files(models_dir)
if gguf_files.empty?
  warn "No GGUF files found in #{models_dir}"
  exit 1
end

puts "Found #{gguf_files.size} GGUF model(s):"
gguf_files.each_with_index do |path, i|
  puts "  #{i + 1}. #{parse_model_name(path)} (#{get_file_size(path)})"
end

selected = []
gguf_files.each_with_index do |path, _i|
  name = parse_model_name(path)
  answer = ask_user("\nInclude #{name} in benchmark? [y/N] ")
  selected << path if %w[y yes].include?(answer)
end

if selected.empty?
  warn 'No models selected. Exiting.'
  exit 0
end

puts "\nRunning benchmark for #{selected.size} model(s)..."

test_configs = [
  { name: 'pp128',  args: '-p 128 -n 0' },
  { name: 'pp2048', args: '-p 2048 -n 0' },
  { name: 'pp8192', args: '-p 8192 -n 0' },
  { name: 'tg256',  args: '-p 0 -n 256' }
]

all_results = []
selected.each do |gguf_path|
  puts "\nBenchmarking: #{parse_model_name(gguf_path)}"
  results = run_bench(gguf_path, test_configs)
  all_results.concat(results)
end

output_file = File.join(File.dirname(__FILE__), 'benchmark_results.csv')
CSV.open(output_file, 'w') do |csv|
  csv << %w[model size test t_s]
  all_results.each { |r| csv << r.values_at(:model, :size, :test, :t_s) }
end
puts "\nResults saved to #{output_file}"