refactor code and support cuda.

Author: wangzhaode

CMakeLists.txt

@@ -12,11 +12,10 @@ link_directories(${CMAKE_CURRENT_LIST_DIR}/libs)
 # source files
 FILE(GLOB SRCS ${CMAKE_CURRENT_LIST_DIR}/src/*.cpp)
 
-# target
-add_executable(chat ${SRCS})
+# compile dynamic lib
+add_library(chat SHARED ${SRCS})
+target_link_libraries(chat MNN MNN_Express)
 
-if (MSVC)
-    target_link_libraries(chat MNN)
-else()
-    target_link_libraries(chat MNN)
-endif()
+# demo target
+add_executable(demo ${CMAKE_CURRENT_LIST_DIR}/demo/main.cpp)
+target_link_libraries(demo chat)

README.md

@@ -1,4 +1,11 @@
 # ChatGLM-MNN
+## Describe
+该模型使用ChatGLM-6B, 将其转换到MNN模型并使用C++进行推理。在实现上做了如下优化：
+
+1. 对其中的词表做了部分删减，删除了模型中未使用的前20000个词；将vocab, embedding, lm_head的大小从130528缩小到130528;
+2. `Embedding`操作调用次数较少，使用`fseek`, `fread`加载的方式降低内存;
+3. `lm_head`操作为`[num, 4096] @ [4096, 130528]`，将其优化为分段实现的矩阵乘`[130528, 4096] @ [4096, 1]`;
+2. 原模型对显存要求较高；将模型按层拆分成28个模型，可以根据用户显存大小动态将计算任务分配给GPU和CPU，充分利用显卡算力;
 
 ## Usage
 ### 1. Compile MNN library
@@ -20,7 +27,7 @@ mkdir build
 cd build
 cmake ..
 make -j8
-./chat
+./demo
 ```
 
 ## Example
@@ -75,4 +82,25 @@ A: 内存泄漏(Memory Leak)是指程序在运行过程中,占用的内存空间
 
 5. 定期清理无用的数据:如果程序中存在大量的无用数据,也会导致内存泄漏。可以通过定期清理这些数据,释放内存空间。
 
-内存泄漏问题的解决需要从多个方面入手,通过不断地调试和优化程序,来找到内存泄漏的根本原因,并有效地解决问题。
+内存泄漏问题的解决需要从多个方面入手,通过不断地调试和优化程序,来找到内存泄漏的根本原因,并有效地解决问题。
+
+---
+Q: 北京有什么特色小吃？
+
+A: 北京是中国历史文化名城,也是中国美食之都之一,有许多特色小吃。以下是一些著名的北京特色小吃:
+
+1. 炸酱面:炸酱面是中国传统面食之一,以黄酱和肉末为主要材料,配以豆瓣酱、黄瓜丝和豆芽等配料,味道鲜美。
+
+2. 烤鸭:烤鸭是北京最著名的美食之一,以薄饼和鸭肉为主要材料,烤制过程中还会加入葱、姜等调料,口感鲜美。
+
+3. 豆汁:豆汁是一种传统的北京小吃,以黄豆为主要原料,配以辣椒油、醋、蒜泥等调料,味道酸甜可口。
+
+4. 羊蝎子:羊蝎子是一道以羊肉和羊肝为主要材料的炖菜,口感鲜美,营养丰富。
+
+5. 糖葫芦:糖葫芦是一种传统的北京小吃,以草莓、山楂等水果为主料,沾上糖浆,口感酸甜可口。
+
+6. 煎饼果子:煎饼果子是一种流行的中式早餐,以薄饼和蛋、肉松、油条等为主要材料,口感酥脆。
+
+7. 驴打滚:驴打滚是一种传统的北京小吃,以糯米粉和豆沙为主要材料,通过卷起来和炸的方式制作,口感香甜。
+
+这只是北京众多特色小吃中的一小部分,北京还有很多其他美食,如北京火锅、北京炸酱面、北京小吃街等等,值得一试。

demo/main.cpp

@@ -0,0 +1,15 @@
+//
+//  chat.cpp
+//
+//  Created by MNN on 2023/03/24.
+//  ZhaodeWang
+//
+
+#include "chat.hpp"
+#include <iostream>
+
+int main(int argc, const char* argv[]) {
+    ChatGLM chatglm;
+    chatglm.chat();
+    return 0;
+}

include/chat.hpp

@@ -0,0 +1,64 @@
+//
+//  chat.hpp
+//
+//  Created by MNN on 2023/03/24.
+//  ZhaodeWang
+//
+
+#ifndef CHAT_hpp
+#define CHAT_hpp
+
+#include <vector>
+#include <memory>
+#include <string>
+#include <unordered_map>
+
+#include <MNN/AutoTime.hpp>
+#include <MNN/expr/Expr.hpp>
+#include <MNN/expr/Module.hpp>
+#include <MNN/expr/MathOp.hpp>
+#include <MNN/expr/NeuralNetWorkOp.hpp>
+
+using namespace MNN;
+using namespace Express;
+
+static constexpr int MASK = 130000;
+static constexpr int gMASK = 130001;
+static constexpr int BOS = 130004;
+static constexpr int EOS = 130005;
+
+static constexpr int VOCAB_SIZE = 130528;
+static constexpr int HIDDEN_SIZE = 4096;
+static constexpr int LAYER_SIZE = 28;
+
+
+class ChatGLM {
+public:
+    // your cuda memory size (G)
+    ChatGLM(float cuda_memory = 0) {
+        init(cuda_memory);
+    }
+    void chat();
+    std::string response(const std::string& input_str, bool debuginfo = false);
+private:
+    void init(float cuda_memory);
+    void loadModel(const char* fileName, bool cuda = false);
+    std::vector<int> tokenizer_encode(std::string input_str);
+    VARP gen_embedding(const std::vector<int>& input_ids);
+    VARP gen_attention_mask(const std::vector<int>& input_ids);
+    VARP gen_position_ids(const std::vector<int>& input_ids);
+    int var_to_token(VARP var);
+    int forward(const std::vector<int>& input_ids);
+private:
+    std::vector<std::string> mWordDecode;
+    std::unordered_map<std::string, int> mWordEncode;
+    // MNN Modules
+    std::shared_ptr<Executor::RuntimeManager> mCPURtmgr;
+    std::shared_ptr<Executor::RuntimeManager> mCUDARtmgr;
+    std::vector<std::shared_ptr<Module>> mModules;
+    std::vector<VARP> mHistoryVars;
+    // mask info
+    int mSeqLen = 0, mContextLen = -1, mMaskIdx = -1;
+};
+
+#endif // CHAT_hpp