Add ChatGLM model type & tokenizer to pybind (#304)

li-plus · Jun 13, 2024 · 1003af9 · 1003af9
1 parent 598b38e
commit 1003af9
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/chatglm_pybind.cpp b/chatglm_pybind.cpp
@@ -49,6 +49,7 @@ PYBIND11_MODULE(_C, m) {
         .value("CHATGLM", ModelType::CHATGLM)
         .value("CHATGLM2", ModelType::CHATGLM2)
         .value("CHATGLM3", ModelType::CHATGLM3)
+        .value("CHATGLM4", ModelType::CHATGLM4)
         .value("BAICHUAN7B", ModelType::BAICHUAN7B)
         .value("BAICHUAN13B", ModelType::BAICHUAN13B)
         .value("INTERNLM", ModelType::INTERNLM);
@@ -143,6 +144,10 @@ PYBIND11_MODULE(_C, m) {
 
     py::class_<ChatGLM3Tokenizer, BaseTokenizer>(m, "ChatGLM3Tokenizer");
 
+    // ===== ChatGLM4 =====
+
+    py::class_<ChatGLM4Tokenizer, BaseTokenizer>(m, "ChatGLM4Tokenizer");
+
     // ===== Baichuan7B/13B =====
 
     py::class_<BaichuanTokenizer, BaseTokenizer>(m, "BaichuanTokenizer");