add get vocab

bindings/python/src/lib.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use pyo3::prelude::*;
 use std::str;
 use rwkv_tokenizer;
@@ -24,6 +25,14 @@ impl WorldTokenizer {
     pub(crate) fn decode(&self, vec: Vec<u16>) -> String {
         return self.tokenizer.decode(vec);
     }
+
+    pub(crate) fn vocab_size(&self) -> usize {
+        return self.tokenizer.vocab_size();
+    }
+
+    pub(crate) fn get_vocab(&self) -> HashMap<&Vec<u8>, usize> {
+        return self.tokenizer.get_vocab();
+    }
 }
 
 

rwkv-tokenizer/src/lib.rs

@@ -1,5 +1,6 @@
 mod trie;
 use std::{str, env};
+use std::collections::HashMap;
 use std::fs::File;
 use std::io::{self, BufRead};
 use std::path::{Path};
@@ -66,6 +67,18 @@ impl WorldTokenizer {
         return str::from_utf8(&*result).unwrap().to_string();
     }
 
+    pub fn vocab_size(&self) -> usize {
+        self.tokens.len()
+    }
+
+    pub fn get_vocab(&self) -> HashMap<&Vec<u8>, usize> {
+        let mut vocabularies: HashMap<&Vec<u8>, usize> = HashMap::new();
+        for (index, value) in self.tokens.iter().enumerate() {
+            vocabularies.insert(value, index);
+        }
+        vocabularies
+    }
+
     fn hex_to_bytes(hex: &str) -> Option<Vec<u8>> {
         let hex = hex.replace("\\x", "");
         if hex.len() % 2 == 0 {