# frozen_string_literal: true
require "net/http"
require "json"
U = URI("http://127.0.0.1:8195/completion")
def one(prompt)
  t = Process.clock_gettime(Process::CLOCK_MONOTONIC)
  res = Net::HTTP.start(U.host, U.port, read_timeout: 600) do |h|
    h.post(U.path, { prompt: prompt, n_predict: 200, temperature: 0, seed: 1, ignore_eos: true, cache_prompt: false }.to_json,
           "Content-Type" => "application/json")
  end
  j = JSON.parse(res.body)
  { wall: Process.clock_gettime(Process::CLOCK_MONOTONIC) - t, gen: j["tokens_predicted"], decode_tps: j.dig("timings", "predicted_per_second") }
end
prompt = "Explain in plain words how a key-value cache speeds up text generation in a transformer. " * 4
one("hi")
ARGV.map(&:to_i).each do |n|
  t = Process.clock_gettime(Process::CLOCK_MONOTONIC)
  rs = Array.new(n) { |i| Thread.new { one("#{i} #{prompt}") } }.map(&:value)
  wall = Process.clock_gettime(Process::CLOCK_MONOTONIC) - t
  puts({ slots: ENV["SLOTS"].to_i, concurrency: n, wall_s: wall.round(2), total_gen_tokens: rs.sum { _1[:gen] },
         aggregate_tps: (rs.sum { _1[:gen] } / wall).round(1), per_req_decode_tps: rs.map { _1[:decode_tps].round(1) },
         max_latency_s: rs.map { _1[:wall] }.max.round(2) }.to_json)
end