Skip to contents

Per-prompt top-k by reward (STABLE ties), retained in original order.

Usage

morie_kamath_rejection_sampling_finetune(
  prompts,
  samples,
  rewards,
  k,
  sft = NULL
)

Arguments

prompts

Prompts.

samples

Per-prompt sample lists.

rewards

Per-prompt reward lists.

k

Kept per prompt.

sft

Optional trainer over the retained pairs.

Value

List with retained, retained_rewards, n_dropped.

Examples

morie_kamath_rejection_sampling_finetune(prompts = c(1, 2, 3, 4, 5, 6, 7, 8),
  samples = c(1, 2, 3, 4, 5, 6, 7, 8), rewards = c(1, 2, 3, 4, 5, 6, 7, 8), k = 5L)
#> $retained
#> $retained[[1]]
#> $retained[[1]][[1]]
#> [1] 1
#> 
#> $retained[[1]][[2]]
#> [1] 1
#> 
#> 
#> $retained[[2]]
#> $retained[[2]][[1]]
#> [1] 2
#> 
#> $retained[[2]][[2]]
#> [1] 2
#> 
#> 
#> $retained[[3]]
#> $retained[[3]][[1]]
#> [1] 3
#> 
#> $retained[[3]][[2]]
#> [1] 3
#> 
#> 
#> $retained[[4]]
#> $retained[[4]][[1]]
#> [1] 4
#> 
#> $retained[[4]][[2]]
#> [1] 4
#> 
#> 
#> $retained[[5]]
#> $retained[[5]][[1]]
#> [1] 5
#> 
#> $retained[[5]][[2]]
#> [1] 5
#> 
#> 
#> $retained[[6]]
#> $retained[[6]][[1]]
#> [1] 6
#> 
#> $retained[[6]][[2]]
#> [1] 6
#> 
#> 
#> $retained[[7]]
#> $retained[[7]][[1]]
#> [1] 7
#> 
#> $retained[[7]][[2]]
#> [1] 7
#> 
#> 
#> $retained[[8]]
#> $retained[[8]][[1]]
#> [1] 8
#> 
#> $retained[[8]][[2]]
#> [1] 8
#> 
#> 
#> 
#> $retained_rewards
#> [1] 1 2 3 4 5 6 7 8
#> 
#> $n_retained
#> [1] 8
#> 
#> $n_dropped
#> [1] 0
#> 
#> $mean_retained_reward
#> [1] 4.5
#> 
#> $k
#> [1] 5
#> 
#> $estimate
#> [1] 8
#> 
#> $n
#> [1] 8
#> 
#> $method
#> [1] "Rejection-sampling fine-tuning (per-prompt top-k)"
#>