From ce32060198b7e2d6a13a9b8e1e1369e3c295ae2a Mon Sep 17 00:00:00 2001
From: Guoteng <32697156+SolenoidWGT@users.noreply.github.com>
Date: Thu, 1 Feb 2024 17:19:51 +0800
Subject: llama : support InternLM2 (#5184)

* support InternLM2 inference
  * add add_space_prefix KV pair
---
 gguf-py/gguf/constants.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

(limited to 'gguf-py/gguf/constants.py')

diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index f5c933a4..ed8e26f8 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -72,6 +72,7 @@ class Keys:
         PAD_ID        = "tokenizer.ggml.padding_token_id"
         ADD_BOS       = "tokenizer.ggml.add_bos_token"
         ADD_EOS       = "tokenizer.ggml.add_eos_token"
+        ADD_PREFIX    = "tokenizer.ggml.add_space_prefix"
         HF_JSON       = "tokenizer.huggingface.json"
         RWKV          = "tokenizer.rwkv.world"
         CHAT_TEMPLATE = "tokenizer.chat_template"
@@ -102,6 +103,7 @@ class MODEL_ARCH(IntEnum):
     PLAMO     = auto()
     CODESHELL = auto()
     ORION     = auto()
+    INTERNLM2  = auto()
 
 
 class MODEL_TENSOR(IntEnum):
@@ -153,6 +155,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.PLAMO:          "plamo",
     MODEL_ARCH.CODESHELL:      "codeshell",
     MODEL_ARCH.ORION:          "orion",
+    MODEL_ARCH.INTERNLM2:      "internlm2",
 }
 
 TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
@@ -446,6 +449,21 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
     ],
+    MODEL_ARCH.INTERNLM2: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_ROT_EMBD,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
     # TODO
 }
 
-- 
cgit v1.2.3