use lru for cached bpe

Project.toml

@@ -8,6 +8,7 @@ Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DoubleArrayTries = "abbaa0e5-f788-499c-92af-c35ff4258c82"
+LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 StructWalk = "31cdf514-beb7-4750-89db-dda9d2eb8d3d"
 TextEncodeBase = "f92c20c0-9f2a-4705-8116-881385faba05"
@@ -16,6 +17,7 @@ Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 [compat]
 DataStructures = "0.18"
 DoubleArrayTries = "0.1"
+LRUCache = "1.6"
 StructWalk = "0.2"
 TextEncodeBase = "0.8"
 julia = "1.6"

src/bpe.jl

@@ -1,3 +1,4 @@
+using LRUCache
 using Unicode
 using DataStructures
 
@@ -28,19 +29,14 @@ struct NoBPE <: AbstractBPE end
 
 Base.show(io::IO, bpe::NoBPE) = print(io, "NoBPE()")
 
-struct CachedBPE{B <: AbstractBPE, D <: AbstractDict{String, Vector{String}}} <: AbstractBPE
+struct CachedBPE{B <: AbstractBPE, D <: AbstractDict{<:AbstractString, Vector{String}}} <: AbstractBPE
     bpe::B
     cache::D
 end
 
-CachedBPE(bpe::AbstractBPE) = CachedBPE(bpe, Dict{String, Vector{String}}())
+CachedBPE(bpe::AbstractBPE) = CachedBPE(bpe, LRU{AbstractString, Vector{String}}(; maxsize = 1000))
 
-function (bpe::CachedBPE)(x)
-    haskey(bpe.cache, x) && return bpe.cache[x]
-    y = bpe.bpe(x)
-    bpe.cache[x] = y
-    return y
-end
+(bpe::CachedBPE)(x) = get!(()->bpe.bpe(x), bpe.cache, x)
 
 Base.show(io::IO, bpe::CachedBPE) = (print(io, "CachedBPE("); show(io, bpe.bpe); print(io, ')'))