add a clear_cache function!

huggingface · Oct 21, 2024 · 6e0175a · 6e0175a
1 parent 5182653
commit 6e0175a
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/bindings/python/src/models.rs b/bindings/python/src/models.rs
@@ -534,6 +534,16 @@ impl PyBPE {
             )?,
         )
     }
+
+    fn clear_cache(self_: PyRef<Self>){
+        let super_ = self_.as_ref();
+        let model = super_.model.read().unwrap();
+        if let ModelWrapper::BPE(ref mo) = *model {
+            mo.clear_cache()
+        } else {
+            unreachable!()
+        };
+    }
 }
 
 /// An implementation of the WordPiece algorithm