Benchmark for Batching LLMs in R

batching-benchmark.R

library(hellmer)
library(tidyllm)
library(ellmer)
library(purrr)
library(tictoc)
library(dplyr)
library(tidyr)
library(ggplot2)
library(patchwork)
library(RColorBrewer)

# prompts -----------------------------------------------------------------

prompts <- list(
  "What is 2+2?",
  "Name one planet.",
  "Is water wet?",
  "What color is the sky?",
  "What is the capital of France?",
  "How many days are in a week?",
  "What is the chemical symbol for gold?",
  "Name a primary color.",
  "What is the largest ocean?",
  "How many continents are there?",
  "Who wrote Hamlet?",
  "What is the boiling point of water?",
  "Name a citrus fruit.",
  "What is the square root of 64?",
  "Who was the first person on the moon?",
  "What is the tallest mountain?",
  "What is the longest river?",
  "Name a mammal that can fly.",
  "What is the largest desert?",
  "How many elements are in the periodic table?",
  "What year did World War II end?",
  "What is photosynthesis?",
  "Who painted the Mona Lisa?",
  "What is the speed of light?",
  "Name a string instrument.",
  "What is the capital of Japan?",
  "What planet is known as the Red Planet?",
  "What is DNA?",
  "How many bones are in the human body?",
  "What is the freezing point of water in Fahrenheit?",
  "Name a prime number.",
  "What is the smallest country by land area?",
  "What gas do plants release during photosynthesis?",
  "Who discovered penicillin?",
  "What is the largest species of shark?",
  "Name a noble gas.",
  "What is the currency of Brazil?",
  "Who invented the telephone?",
  "What is the closest star to Earth?",
  "What is the capital of Australia?",
  "Name a bird that cannot fly.",
  "What is the main ingredient in guacamole?",
  "Who is the author of '1984'?",
  "What is the hardest natural substance?",
  "Name a wind instrument.",
  "What is the capital of Egypt?",
  "What is the largest planet in our solar system?",
  "What is the chemical formula for water?",
  "How many teeth does an adult human have?",
  "What is the national flower of Japan?",
  "Name a type of cloud.",
  "What is the largest organ in the human body?",
  "Who composed the Four Seasons?",
  "What is the capital of Canada?",
  "What element has the symbol 'Fe'?",
  "Name a rodent.",
  "What is the main language spoken in Brazil?",
  "Who discovered gravity?",
  "What is the smallest prime number?",
  "Name a flightless bird.",
  "What is the largest species of bear?",
  "What is the capital of South Korea?",
  "What planet has the most moons?",
  "What is the chemical symbol for silver?",
  "How many sides does a hexagon have?",
  "What is the deepest ocean trench?",
  "Name a constellation.",
  "What is the official language of China?",
  "Who wrote 'Pride and Prejudice'?",
  "What is the square root of 144?",
  "Name a nocturnal animal.",
  "What is the capital of Mexico?",
  "What element has the atomic number 1?",
  "What is the largest land animal?",
  "Name a percussion instrument.",
  "What is the boiling point of water in Fahrenheit?",
  "Who painted 'Starry Night'?",
  "What is the fastest land animal?",
  "What is the capital of Italy?",
  "What is the chemical symbol for potassium?",
  "How many planets are in our solar system?",
  "Name a cold-blooded animal.",
  "What is the largest bird of prey?",
  "What is the capital of Russia?",
  "What gas do humans breathe out?",
  "Who invented the light bulb?",
  "What is the smallest ocean?",
  "Name a monotreme.",
  "What is the capital of Argentina?",
  "What is the chemical symbol for sodium?",
  "How many faces does a cube have?",
  "What is the study of fossils called?",
  "Name a poisonous snake.",
  "What is the capital of Spain?",
  "What is the third planet from the sun?",
  "What is the chemical formula for table salt?",
  "What is the average lifespan of a honey bee?",
  "Name three cloud types.",
  "How do earthquakes occur?",
  "What is the difference between weather and climate?",
  "How does a vaccine work?",
  "What is the purpose of sleep?",
  "How many chromosomes do humans have?",
  "Explain how rainbows form.",
  "What causes the northern lights?",
  "How do plants respond to light?",
  "What are black holes?",
  "Explain how sound travels.",
  "What is the function of red blood cells?",
  "How are mountains formed?",
  "What's the difference between a virus and bacteria?",
  "How does GPS work?",
  "What causes seasons on Earth?",
  "How do airplanes fly?",
  "What is machine learning?",
  "How does a refrigerator work?",
  "What are the major types of renewable energy?",
  "How do coral reefs form?",
  "What is the water cycle?",
  "How do volcanoes erupt?",
  "What is the role of bees in ecosystems?",
  "How does a computer process information?",
  "What causes thunder and lightning?",
  "How do humans digest food?",
  "What is biodiversity and why is it important?",
  "How do vaccines create immunity?",
  "What are the main components of soil?",
  "How do stars form?",
  "What causes ocean tides?",
  "How do migrating birds navigate?",
  "What is cellular respiration?",
  "How does a camera work?",
  "What causes earthquakes?",
  "How do nuclear power plants generate electricity?",
  "What is the greenhouse effect?",
  "How do trees communicate?",
  "What is artificial intelligence?",
  "How do electric cars work?",
  "What causes acid rain?",
  "How do we measure the age of fossils?",
  "What is blockchain technology?",
  "How do solar panels convert sunlight to electricity?",
  "What happens during a solar eclipse?",
  "How do antibiotics work?",
  "What causes the aurora borealis?",
  "How do whales communicate?",
  "What is quantum computing?",
  "How do clouds form?",
  "What causes galaxies to form?",
  "How does the human immune system work?",
  "What is the impact of plastic pollution on marine life?",
  "How do wind turbines generate electricity?",
  "What causes tsunamis?",
  "How do touchscreens work?",
  "What is natural selection?",
  "How do MRI machines create images of the body?",
  "What causes climate change?",
  "How do bees make honey?",
  "What is gene editing?",
  "How do batteries store and release energy?",
  "What causes desertification?",
  "How do plants adapt to different environments?",
  "What is the Internet of Things?",
  "How do submarines work?",
  "What causes sea level rise?",
  "How do chameleons change color?",
  "What is nuclear fusion?",
  "How do hearing aids work?",
  "What causes coral bleaching?",
  "How do self-driving cars navigate?",
  "What is cognitive bias?",
  "How do 3D printers work?",
  "What causes sinkholes?",
  "How do plants respond to drought?",
  "What is virtual reality?",
  "How do fireflies produce light?",
  "What causes hurricanes to form?",
  "How do humans perceive color?",
  "What is deep learning?",
  "How do geothermal energy systems work?",
  "What causes mental fatigue?",
  "How do enzymes catalyze reactions?",
  "What is CRISPR technology?",
  "How do microwave ovens heat food?",
  "What causes landslides?",
  "How do space telescopes work?",
  "What is neural plasticity?",
  "How do smoke detectors function?",
  "What causes wildfires to spread?",
  "How do plants communicate with each other?",
  "What is augmented reality?",
  "How do noise-cancelling headphones work?",
  "What causes ocean acidification?",
  "How do birds fly in formation?",
  "How do fungi contribute to ecosystem health?",
  "What causes déjà vu experiences?",
  "How do quantum computers differ from classical computers?",
  "What is the relationship between gut bacteria and mental health?",
  "How do smart materials respond to environmental changes?",
  "What are the ethical implications of advanced AI systems?"
)

# benchmark ---------------------------------------------------------------

run_hellmer_sequential <- function(prompts_list) {
  chat <- chat_sequential(
    chat_openai,
    system_prompt = "Reply concisely in one sentence"
  )
  result <- chat$batch(prompts_list)
}

run_hellmer_parallel <- function(prompts_list) {
  chat <- chat_future(
    chat_openai,
    system_prompt = "Reply concisely in one sentence",
    workers = 30L,
    chunk_size = length(prompts)
  )
  result <- chat$batch(prompts_list)
}

run_ellmer_parallel <- function(prompts_list) {
  cht <- ellmer::chat_openai(system_prompt = "Reply concisely in one sentence")
  responses <- cht$chat_parallel(prompts_list)
}

run_ellmer_sequential <- function(prompts_list) {
  cht <- ellmer::chat_openai(system_prompt = "Reply concisely in one sentence", echo = "none")
  responses <- purrr::map(
    prompts_list,
    ~ cht$clone()$chat(.x, echo = "none")
  )
}

n_iterations <- 5
results <- data.frame(
  iteration = numeric(),
  method = character(),
  time_seconds = numeric()
)

for (i in 1:n_iterations) {
  test_functions <- list(
    hellmer_sequential = function() run_hellmer_sequential(prompts),
    hellmer_parallel = function() run_hellmer_parallel(prompts),
    ellmer_parallel = function() run_ellmer_parallel(prompts),
    ellmer_sequential = function() run_ellmer_sequential(prompts)
  )

  test_order <- sample(names(test_functions))
  iteration_results <- data.frame()

  for (test_name in test_order) {
    tic()
    test_functions[[test_name]]()
    time_result <- toc(quiet = TRUE)

    iteration_results <- bind_rows(
      iteration_results,
      data.frame(
        iteration = i,
        method = test_name,
        time_seconds = time_result$toc - time_result$tic
      )
    )
  }

  results <- bind_rows(results, iteration_results)
}

summary_stats <- results |>
  group_by(method) |>
  summarise(time = round(mean(time_seconds), 0))

print(summary_stats)

# visualization -----------------------------------------------------------

plot_data <- results |>
  tidyr::separate(method, into = c("library", "execution_type"), sep = "_")

parallel_data <- plot_data |>
  dplyr::filter(execution_type == "parallel")

sequential_data <- plot_data |>
  dplyr::filter(execution_type == "sequential")

p1 <- parallel_data |>
  ggplot2::ggplot(aes(x = library, y = time_seconds, fill = library)) +
  ggplot2::geom_violin(alpha = 0.7, width = 0.8) +
  ggplot2::geom_point(
    aes(color = library),
    position = position_jitter(width = 0.1, seed = 123),
    size = 3,
    alpha = 0.8
  ) +
  ggplot2::scale_fill_brewer(palette = "Set2") +
  ggplot2::scale_color_brewer(palette = "Set2") +
  ggplot2::scale_y_continuous(
    breaks = function(limits) seq(0, ceiling(max(limits)), by = 1)
  ) +
  ggplot2::labs(subtitle = "Parallel") +
  ggplot2::theme_minimal() +
  ggplot2::theme(
    legend.position = "none",
    axis.title.x = element_blank()
  )

p2 <- sequential_data |>
  ggplot2::ggplot(aes(x = library, y = time_seconds, fill = library)) +
  ggplot2::geom_violin(alpha = 0.7, width = 0.8) +
  ggplot2::geom_point(
    aes(color = library),
    position = position_jitter(width = 0.1, seed = 123),
    size = 3,
    alpha = 0.8
  ) +
  ggplot2::scale_fill_brewer(palette = "Set2") +
  ggplot2::scale_color_brewer(palette = "Set2") +
  ggplot2::scale_y_continuous(
    breaks = function(limits) seq(0, ceiling(max(limits) / 5) * 5, by = 5)
  ) +
  ggplot2::labs(subtitle = "Sequential") +
  ggplot2::theme_minimal() +
  ggplot2::theme(
    legend.position = "none",
    axis.title.y = element_blank(),
    axis.title.x = element_blank()
  )

combined_plot <- p1 + p2 +
  plot_layout(ncol = 2) +
  plot_annotation(
    title = "Real-Time Batch Performance",
    caption = "OpenAI, random order, 200 prompts x 5 runs"
  ) &
  theme(
    plot.title = element_text(size = 14, hjust = 0.5),
    strip.text = element_text(size = 12, face = "bold"),
    axis.title = element_text(size = 11, face = "bold"),
    axis.text = element_text(size = 10)
  ) &
  ggplot2::ylab("Seconds")

print(combined_plot)