Un-deprecate sampler interface

MurrellGroup · Nov 28, 2024 · 7d4f50b · 7d4f50b · AntonOresten · Nov 28, 2024
1 parent b2a4c93
commit 7d4f50b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 12 deletions.
diff --git a/src/LogitSamplers.jl b/src/LogitSamplers.jl
@@ -16,7 +16,7 @@ export Top_pk, Top_p, Top_k
 export Min_p
 export Top_nσ
 
-include("deprecated.jl")
+include("samplers.jl")
 export argmax_sampler, top_pk_sampler, min_p_sampler, top_nσ_sampler
 
 end
diff --git a/src/deprecated.jl b/src/deprecated.jl
diff --git a/src/samplers.jl b/src/samplers.jl
@@ -0,0 +1,11 @@
+argmax_sampler(logits; device=identity) = logits |> device |> Top_k(1) |> logitsample
+argmax_sampler(; kwargs...) = logits -> argmax_sampler(logits; kwargs...)
+
+top_pk_sampler(logits; p = 0.5f0, k = 5, device = identity) = logits |> device |> Top_pk(p, k) |> logitsample
+top_pk_sampler(; kwargs...) = logits -> top_pk_sampler(logits; kwargs...)
+
+min_p_sampler(logits; pbase = 0.5f0, device = identity) = logits |> device |> Min_p(pbase) |> logitsample
+min_p_sampler(; kwargs...) = logits -> min_p_sampler(logits; kwargs...)
+
+top_nσ_sampler(logits; temperature = 1.0f0, n = 1.0f0, device = identity) = logits |> device |> Temperature(temperature) |> Top_nσ(n) |> logitsample
+top_nσ_sampler(; kwargs...) = logits -> top_nσ_sampler(logits; kwargs...)