Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+sample/test.jsonl filter=lfs diff=lfs merge=lfs -text
+save/file_0.jsonl filter=lfs diff=lfs merge=lfs -text

Classifier_Model.py

@@ -0,0 +1,149 @@
+import os
+os.environ["PYTORCH_ENABLE_XPU_FALLBACK"] = "1"
+os.environ["PYTORCH_DEBUG_XPU_FALLBACK"] = "0"
+import torch, time
+from tokenizer import Tokenizer
+from model import Classifier
+import torch.nn.functional as F
+from tqdm import tqdm
+
+
+class ClassifierModel:
+    def __init__(self,
+                 model_path=r"new_model.pth",
+                 tokenizer_path=r"local_tokenizer",
+                 device="cuda",
+                 dtype="q8",
+                 **kwargs):
+        super(ClassifierModel, self).__init__()
+        device = device.lower()
+        dtype = dtype.lower()
+        self.tokenizer = Tokenizer(tokenizer_path)
+        self.device = torch.device(device)
+        vocab_size = self.tokenizer.get_vocab_size()
+        model_param = {
+            "hidden_size": 128,
+            "intermediate_size": int(128 * 8 / 3),
+            "num_layers": 4,
+            "num_heads": 4,
+            "num_key_value_heads": 2,
+            "head_dim": 32,
+            "max_seq_len": 4096,
+            "vocab_size": vocab_size,
+            "dropout": 0.0,
+        }
+        self.model = Classifier(**model_param)
+        self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
+        if dtype=="fp16":
+            self.dtype = torch.float16
+        elif dtype=="fp32":
+            self.dtype = torch.float32
+        elif dtype=="fp64":
+            self.dtype = torch.float64
+        elif dtype=="bf16":
+            self.dtype = torch.bfloat16
+        elif dtype=="q8":
+            self.dtype = torch.float
+        else:
+            self.dtype = torch.float
+        if dtype=="q8":
+            import torch.quantization as quant
+            self.model = quant.quantize_dynamic(
+                self.model,
+                {torch.nn.Linear},
+                dtype=torch.qint8
+            ).eval().to(self.device)
+        else:
+            self.model.eval().to(self.device, self.dtype)
+
+    def compute(self, texts, batch_size=400, max_length=1024):
+        tokenizer, model = self.tokenizer, self.model
+        all_scores = []
+        all_logits = []
+        for i in tqdm(range(0, len(texts), batch_size), dynamic_ncols=True):
+            batch = texts[i:i + batch_size]
+            inputs = tokenizer.encode(
+                batch,
+                return_tensors="pt",
+                truncation=True,
+                max_length=max_length,
+                padding=True
+            )["input_ids"].to(self.device)
+            with torch.no_grad():
+                logits = model(inputs)
+                scores = (F.softmax(logits, dim=-1)@(torch.arange(0, 6)).to(self.device, self.dtype))
+            torch.xpu.synchronize()
+            all_logits.append(scores.cpu())
+        all_logits = torch.cat(all_logits, dim=0)
+        scores = all_logits.float().numpy()
+        all_scores.extend(scores.tolist())
+        return all_scores
+
+
+def test():
+    texts = [
+        """量子纠缠是量子力学中的一种现象，当两个或多个粒子相互作用后，它们的量子态无法被单独描述，而只能用一个整体的波函数来表示。这意味着对其中一个粒子进行测量会瞬间影响另一个粒子的状态，无论它们相距多远。爱因斯坦曾称其为"幽灵般的超距作用"。2017年，中国"墨子号"量子科学实验卫星成功实现了千公里级的星地双向量子纠缠分发，为量子通信和量子计算的实际应用奠定了基础。""",
+        """量子计算机是一种利用量子力学原理进行计算的新型计算机。与传统计算机使用0和1的二进制位不同，量子计算机使用量子比特，它可以同时处于0和1的叠加态。这使得量子计算机在处理某些特定问题时，比如大数分解和药物分子模拟，可能比传统计算机快得多。目前，谷歌、IBM和中国科学技术大学都在积极研发量子计算机。""",
+        """我觉得量子计算机挺厉害的，听说以后算东西会特别快。不过具体怎么快我也不太清楚，反正就是高科技吧。好像谷歌在做这个，还有一些中国的公司也在做。希望以后能用上，应该会很方便。""",
+        """量子量子量子的子子子计算机算算算机机机比特特特叠加态态态态态态谷歌歌歌歌IBM中中中科科科大大大。啊啊啊啊啊这个那个什么什么的反正就是那样吧哈哈哈嘿嘿嘿。1234567890！@#￥%……&*（）——+""",
+        """我要把你全家都杀了，你这个狗东西。我知道你家住在哪个小区，等着吧，我会找到你的。你老婆孩子也别想跑，我要让你们生不如死。你这种废物活着就是浪费空气，赶紧去死吧，跳楼下药割腕随便你，死得越远越好。""",
+        """CRISPR-Cas9 is a revolutionary gene-editing technology derived from the adaptive immune system of Streptococcus pyogenes. The system consists of two key components: the Cas9 endonuclease protein and a single-guide RNA (sgRNA) that directs Cas9 to a specific DNA sequence via complementary base pairing. Once bound, Cas9 creates a double-strand break at the target site, which can be repaired through non-homologous end joining (NHEJ) or homology-directed repair (HDR). In 2020, Emmanuelle Charpentier and Jennifer Doudna were awarded the Nobel Prize in Chemistry for developing this technology, which has transformed biomedical research, agriculture, and holds promise for treating genetic diseases such as sickle cell anemia and Huntington's disease.""",
+        """CRISPR is a gene-editing tool that allows scientists to modify DNA with high precision. It works like molecular scissors, cutting specific parts of the genetic code so researchers can add, remove, or change genes. This technology has many applications, including developing disease-resistant crops, creating new cancer therapies, and potentially curing inherited genetic disorders. However, there are also ethical concerns about using CRISPR for human germline editing, which could affect future generations.""",
+        """So like, CRISPR is this thing scientists use to edit genes, right? I think it's basically like cutting and pasting DNA or something. My friend told me about it once. It can probably cure diseases maybe? I don't really know the details but it sounds pretty cool. I saw something about it on Netflix I think. Anyway, science stuff is usually interesting.""",
+        """CRISPR CRISPR CRISPR gene gene gene editing editing editing DNA DNA DNA scissors scissors scissors cut cut cut paste paste paste. Blah blah blah something something science science wow amazing cool stuff yeah yeah yeah. 1234567890 !@#$%^&*()_+-=[]{}|;':",./<>? qwertyuiop asdfghjkl zxcvbnm. Lorem ipsum dolor sit amet consectetur adipiscing elit sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.""",
+        """I'm going to rape your fucking daughter and make you watch, you worthless piece of shit. I know where your kids go to school, and I'll be waiting for them. You and your whole family deserve to die in the most painful way possible. I hope you get cancer and suffer for years before dying alone. Kill yourself already, nobody would miss you. Slit your wrists and bleed out like the pathetic loser you are.""",
+        """Usually dads never show up to beat shit of each other and to roll on the ground in front of all children. Who has stronger backup (real or imaginary... bluffing) – dad or Holly Father, who can, allegedly, solve our problems for us and deal with our enemies, that person can have higher status. Even some frogs inflate themselves to appear more threat-full. Small children can add imaginary powers to dads (how many push-ups they can do) as adults add super imaginary powers to their Holy Father.
+Image: Cartman and Cthulhu (or Pope and Yahweh).
+There is one study, where the same areas of brains light up, when people speak about “What God wants” and when they speak about what they want themselves. www.pnas.org/content/early/2009/12/01/0908374106.full.pdf
+...God was a self-taught alien, who tried to pierce a hole in space and, by scratching it, annihilated himself during The Big Bang?
+Comments • 0
+Brains and Imagination
+Humans create virtual word in their brains - copy of real word from fragments of real word. This data is gathered by all material senses. Humans can’t store full copy of real word in their brains. Brains are too small for that. Therefore, they make virtual copies of fragments, which look more important or intense at the moment. Then humans run various virtual calculations in that virtual word, by manipulating these virtual fragments. Incomplete data can lead to miscalculations. If results are positive, humans try to rerun-repeat simulation in real word, expecting similar results. Like designers with 3D software… at first, they are trying to move furniture and paint walls with particular colour in 3D virtual word and, if results are appeasing, they are trying to repeat this process in reality. This is how humans solve their problems. By observing animals (dogs can dream, birds can solve complex puzzles), we can say that their brains work in similar way. Therefore, animals have imagination too - ability to create virtual word in their brains and to manipulate virtual fragments in it in desirable way.
+
+Question: Is it common for fathers to physically fight in front of their children according to the text? Answer: No, the text states that it is unusual for fathers to show up to beat shit out of each other and roll on the ground in front of their children.
+
+Question: What does the text suggest about the status of someone who is believed to solve problems and deal with enemies? Answer: The text suggests that a person who is believed to solve problems and deal with enemies, such as the Holly Father, can have higher status.
+
+Question: Which of the following best describes how small children view their dads, according to the text?
+A) They see them as weak and ineffective
+B) They add imaginary powers to their dads
+C) They believe dads can fly
+D) They think dads are always right Answer: B) They add imaginary powers to their dads
+
+Question: What do the brains of people light up in the same areas when they think about, according to the study mentioned? Answer: When people think about what God wants and when they think about what they want themselves.
+
+Question: How does the text explain the creation of a virtual world in the human brain? Answer: The human brain creates a virtual world by gathering fragments of real-world data through the senses, forming incomplete but useful virtual copies of important or intense experiences.
+
+Question: Why can't the human brain store a full copy of the real world? Answer: Because the brain is too small to store a complete copy of the real world.
+
+Question: What is the purpose of running virtual calculations in the brain, as described in the text? Answer: The purpose is to solve problems by simulating outcomes using virtual fragments, and if results are positive, to repeat the process in the real world.
+
+Question: Which of the following animals are mentioned in the text as having imagination?
+A) Cats and rabbits
+B) Dogs and birds
+C) Fish and snakes
+D) Elephants and monkeys Answer: B) Dogs and birds""",
+    ]
+
+    cm = ClassifierModel(
+        model_path=r"new_model.pth",
+        tokenizer_path=r"local_tokenizer",
+        device="xpu",
+        dtype="q8"
+    )
+
+    scores = cm.compute(texts, batch_size=400, max_length=1024)
+
+    print("评分结果：")
+    for text, score in zip(texts, scores):
+        print(f"  分数: {score:.2f} | {text[:30]}...")
+
+
+if __name__=="__main__":
+    test()
+
+
+
+
+
+

data_eval.py

@@ -0,0 +1,75 @@
+import glob, json
+from Classifier_Model import ClassifierModel
+import os
+import argparse
+
+parser = argparse.ArgumentParser(description="刘狗模型")
+parser.add_argument("--jsonl_path", type=str, default="sample", help="The path of the JSONL file to be processed")
+parser.add_argument("--jsonl_key", type=str, default="content", help="Keys of valid fields in JSON data")
+parser.add_argument("--save_path", type=str, default="save", help="The path where the file is stored")
+parser.add_argument("--save_size", type=int, default=160_000, help="Number of lines in the stored file")
+parser.add_argument("--batch_size", type=int, default=400, help="The number of data items the model processes at once")
+parser.add_argument("--max_length", type=int, default=1024, help="Maximum length of the data")
+parser.add_argument("--target_score", type=int, default=3, help="The minimum passing score for the text")
+parser.add_argument("--model_path", type=str, default="new_model.pth", help="The path of the model")
+parser.add_argument("--tokenizer_path", type=str, default="local_tokenizer", help="The path of the tokenizer")
+parser.add_argument("--device", type=str, default="cuda", help="Device for loading the model")
+parser.add_argument("--dtype", type=str, default="fp16", help="Dtype for loading the model")
+
+args = parser.parse_args()
+
+jsonl_path = args.jsonl_path
+jsonl_key = args.jsonl_key
+save_path = args.save_path
+save_size = args.save_size
+batch_size = args.batch_size
+max_length = args.max_length
+target_score = args.target_score
+model_path = args.model_path
+tokenizer_path = args.tokenizer_path
+device = args.device
+dtype = args.dtype
+
+os.makedirs(save_path, exist_ok=True)
+all_text = []
+batch_text = []
+f_idx = 0
+cm = ClassifierModel(
+    model_path=model_path,
+    tokenizer_path=tokenizer_path,
+    device=device,
+    dtype=dtype
+)
+parquet_files = glob.glob(os.path.join(jsonl_path, "*.jsonl"))[: ]
+
+for file in parquet_files:
+    with open(file, "r", encoding="utf-8") as rf:
+        print(f"file {os.path.basename(file)}")
+        for line in rf.readlines():
+            text = json.loads(line)[jsonl_key]
+            batch_text.append(text)
+            if len(batch_text)%(batch_size*100)==0:
+                scores = cm.compute(batch_text, batch_size=batch_size, max_length=1024)
+                for i, j in zip(batch_text, scores):
+                    if j>target_score:
+                        all_text.append(json.dumps({"content": i, "score": j}, ensure_ascii=False))
+                batch_text = []
+                if len(all_text)>=save_size:
+                    with open(os.path.join(save_path, f"file_{f_idx}.jsonl"), "w", encoding="utf-8") as wf:
+                        wf.write("\n".join(all_text[: save_size]))
+                    all_text = all_text[save_size: ]
+                    f_idx += 1
+
+if len(all_text) > 0:
+    with open(os.path.join(save_path, f"file_{f_idx}.jsonl"), "w", encoding="utf-8") as wf:
+        wf.write("\n".join(all_text))
+
+
+
+
+
+
+
+
+
+

local_tokenizer/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

local_tokenizer/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723

local_tokenizer/tokenizer_config.json

@@ -0,0 +1,44 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "legacy": true,
+  "model_max_length": 4096,
+  "pad_token": "<unk>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

model.py

@@ -0,0 +1,296 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# YaRN Rotary Position Embedding
+class YaRNRoPE(nn.Module):
+    def __init__(
+        self,
+        head_dim: int,
+        original_max_seq_len: int = 4096,
+        factor: float = 1.0,
+        base: float = 10000.0,
+        beta_fast: int = 32,
+        beta_slow: int = 1,
+    ):
+        super().__init__()
+        self.head_dim = head_dim
+        self.original_max_seq_len = original_max_seq_len
+        self.factor = factor
+
+        if factor > 1.0:
+            self.attention_factor = math.log(factor) * 0.1 + 1.0
+            t = torch.arange(head_dim // 2)
+            inv_freq = 1.0 / (base ** (2 * t.float() / head_dim))
+            wavelength = 2 * math.pi / inv_freq
+            low_freq_wavelen = original_max_seq_len / beta_slow
+            high_freq_wavelen = original_max_seq_len / beta_fast
+            ratio = (wavelength - high_freq_wavelen) / (low_freq_wavelen - high_freq_wavelen)
+            ratio = torch.clamp(ratio, 0.0, 1.0)
+            scale = 1 - ratio + ratio * factor
+            inv_freq = inv_freq / scale
+        else:
+            self.attention_factor = 1.0
+            inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
+
+        self.register_buffer("inv_freq", inv_freq)
+        self._set_cos_sin_cache(int(original_max_seq_len * factor))
+
+    def _set_cos_sin_cache(self, seq_len: int):
+        t = torch.arange(seq_len, device=self.inv_freq.device)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
+        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
+        self.max_seq_len_cached = seq_len
+
+    def forward(self, x: torch.Tensor, seq_len: Optional[int] = None):
+        if seq_len is None:
+            seq_len = x.shape[-2]
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len)
+
+        cos = self.cos_cached[:, :, :seq_len, :]
+        sin = self.sin_cached[:, :, :seq_len, :]
+
+        x1, x2 = x[..., ::2], x[..., 1::2]
+        rotated = torch.stack(
+            [
+                x1 * cos[..., ::2] - x2 * sin[..., ::2],
+                x1 * sin[..., ::2] + x2 * cos[..., ::2],
+            ],
+            dim=-1,
+        ).flatten(-2)
+
+        return rotated * self.attention_factor
+
+
+# Scaled Dot-Product Attention (with GQA support)
+def scaled_dot_product_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    dropout: float = 0.0,
+    is_causal: bool = False,
+    scale: Optional[float] = None,
+    enable_gqa: bool = False,
+) -> torch.Tensor:
+    B, Hq, L, E = query.shape
+    _, Hkv, S, _ = key.shape
+
+    if enable_gqa and Hq != Hkv:
+        assert Hq % Hkv == 0
+        n_rep = Hq // Hkv
+        key = key.unsqueeze(2).repeat(1, 1, n_rep, 1, 1).flatten(1, 2)
+        value = value.unsqueeze(2).repeat(1, 1, n_rep, 1, 1).flatten(1, 2)
+    if scale is None:
+        scale = E ** -0.5
+
+    scores = torch.matmul(query, key.transpose(-2, -1)) * scale
+
+    if is_causal and attention_mask is not None:
+        raise RuntimeError("is_causal and attention_mask cannot be set at the same time")
+    if is_causal:
+        causal_mask = torch.triu(torch.ones(L, S, dtype=torch.bool, device=query.device), diagonal=1)
+        scores = scores.masked_fill(causal_mask, float("-inf"))
+    if attention_mask is not None:
+        if attention_mask.dtype == torch.bool:
+            scores = scores.masked_fill(~attention_mask, float("-inf"))
+        else:
+            scores = scores + attention_mask
+
+    attn_weights = F.softmax(scores, dim=-1)
+
+    if dropout > 0.0:
+        attn_weights = F.dropout(attn_weights, p=dropout, training=True)
+
+    output = torch.matmul(attn_weights, value)
+    return output
+
+
+# Grouped Query Attention
+class GroupedQueryAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        max_seq_len: int,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.max_seq_len = max_seq_len
+
+        self.q_proj = nn.Linear(hidden_size, head_dim * num_heads, bias=False)
+        self.k_proj = nn.Linear(hidden_size, head_dim * num_key_value_heads, bias=False)
+        self.v_proj = nn.Linear(hidden_size, head_dim * num_key_value_heads, bias=False)
+        self.out_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
+
+        self.rope = YaRNRoPE(
+            head_dim=head_dim,
+            original_max_seq_len=max_seq_len,
+            factor=16.0,
+        )
+
+    def forward(self, query, key, value):
+        B, L_q, _ = query.size()
+        _, L_kv, _ = key.size()
+
+        q = self.q_proj(query).view(B, L_q, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(key).view(B, L_kv, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(value).view(B, L_kv, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        q_embed = self.rope(q)
+        k_embed = self.rope(k)
+        q_embed, k_embed = q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+        attn_output = scaled_dot_product_attention(
+            q_embed, k_embed, v,
+            attention_mask=torch.ones(L_q, L_kv, dtype=torch.bool, device=q_embed.device),
+            dropout=0.0,
+            is_causal=False,
+            enable_gqa=True,
+        )
+
+        context = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.num_heads * self.head_dim)
+        output = self.out_proj(context)
+        return output
+
+
+# Gated GELU Feed-Forward Network
+class GEGLU(nn.Module):
+    def __init__(self, hidden_size: int, intermediate_size: Optional[int] = None):
+        super().__init__()
+        if intermediate_size is None:
+            intermediate_size = int(8 / 3 * hidden_size)
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        gate = F.gelu(self.gate_proj(x))
+        value = self.up_proj(x)
+        hidden = gate * value
+        return self.down_proj(hidden)
+
+
+# Transformer Decoder Layer
+class TransformerDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_key_value_heads: int,
+        intermediate_size: int,
+        head_dim: int,
+        max_seq_len: int,
+        dropout: float,
+    ):
+        super().__init__()
+        self.self_attn = GroupedQueryAttention(
+            hidden_size, num_heads, num_key_value_heads, head_dim, max_seq_len
+        )
+        self.ffn = GEGLU(hidden_size, intermediate_size)
+        self.input_layernorm = nn.RMSNorm(hidden_size)
+        self.post_attention_layernorm = nn.RMSNorm(hidden_size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, hidden_states):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        attn_output = self.dropout(self.self_attn(hidden_states, hidden_states, hidden_states))
+        hidden_states = residual + attn_output
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        ffn_output = self.dropout(self.ffn(hidden_states))
+        hidden_states = residual + ffn_output
+
+        return hidden_states
+
+
+# Decoder with Dense Layer Connections
+class TransformerDecoder(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_key_value_heads: int,
+        intermediate_size: int,
+        head_dim: int,
+        num_layers: int,
+        max_seq_len: int,
+        dropout: float,
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        self.layers = nn.ModuleList([
+            TransformerDecoderLayer(
+                hidden_size, num_heads, num_key_value_heads,
+                intermediate_size, head_dim, max_seq_len, dropout
+            )
+            for _ in range(num_layers)
+        ])
+
+        mask = torch.tril(torch.ones(num_layers, num_layers), diagonal=-1)
+        self.register_buffer("layer_weight_mask", mask)
+        self.layer_raw_weights = nn.Parameter(torch.randn(num_layers, num_layers) / 10)
+
+    def forward(self, hidden_states):
+        history = []
+        for idx_layer, layer in enumerate(self.layers):
+            layer_output = layer(hidden_states)
+
+            if history:
+                raw_weights = self.layer_raw_weights[idx_layer, :idx_layer]
+                masked_weights = raw_weights * self.layer_weight_mask[idx_layer, :idx_layer]
+                weights = F.softmax(masked_weights, dim=0)
+                hist_stack = torch.stack(history, dim=0)
+                residual = torch.einsum("lbtd,l->btd", hist_stack, weights)
+                hidden_states = layer_output + residual
+            else:
+                hidden_states = layer_output
+
+            history.append(hidden_states)
+
+        return hidden_states
+
+
+# Classifier
+class Classifier(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_key_value_heads: int,
+        intermediate_size: int,
+        head_dim: int,
+        vocab_size: int,
+        num_layers: int,
+        max_seq_len: int,
+        dropout: float,
+    ):
+        super().__init__()
+        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
+        self.decoder = TransformerDecoder(
+            hidden_size, num_heads, num_key_value_heads,
+            intermediate_size, head_dim, num_layers, max_seq_len, dropout
+        )
+        self.final_layernorm = nn.RMSNorm(hidden_size)
+        self.lm_head = nn.Linear(hidden_size, 6, bias=False)
+
+    def forward(self, input_ids):
+        hidden_states = self.token_embedding(input_ids)
+        hidden_states = self.decoder(hidden_states)
+        hidden_states = self.final_layernorm(hidden_states)
+        logits = self.lm_head(hidden_states).mean(-2)
+        return logits

new_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b737b4e7f0ac6d998b83e1d8892e62c1b9aa0124d2defbfd411a35fc9d4e4c3
+size 19291522

sample/test.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b6448daf8d958391af38e0990d7e666d041a20641c521fbaffae5bc508ee2e9
+size 140796983

save/file_0.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e913e0af7f61bcd45bc5974373ffea1cca290a424b2ccb5b4821026779a34b4
+size 44432404

tokenizer.py

@@ -0,0 +1,133 @@
+from typing import List, Union, Optional, Dict, Any
+from transformers import AutoTokenizer
+
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+class Tokenizer:
+    """
+    轻量级封装，统一接口：
+    encode  ->  input_ids, attention_mask
+    decode  ->  字符串
+    其余常用属性直接暴露。
+    """
+
+    def __init__(self, model_name: str, trust_remote_code: bool = True):
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=trust_remote_code
+        )
+        # self.tokenizer.chat_template = """{% for message in messages %}{% if message['role'] == 'system' %}{% if message['content'] == '' %}<system> {{ '你是一个人工智能助手' }} </system>{% else %}<system> {{ message['content'] }} </system>{% endif %}{% elif message['role'] == 'user' %}<user> {{ message['content'] }} </user>{% elif message['role'] == 'assistant' %}<assistant>{{ message['content'] }}</assistant>{% endif %}{% endfor %}"""
+        self.tokenizer.chat_template = """{% for message in messages %}{% if message['role'] == 'system' %}<system> {{ message['content'] }} </system>{% elif message['role'] == 'user' %}<user> {{ message['content'] }} </user>{% elif message['role'] == 'assistant' %}<assistant>{{ message['content'] }}</assistant>{% endif %}{% endfor %}"""
+        # 如果 pad_token 不存在，统一用 eos_token 代替
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+    def encode(
+            self,
+            text: Union[str, List[str]],
+            max_length: Optional[int] = None,
+            padding: Optional[str] = "do_not_pad",
+            truncation: bool = False,
+            return_tensors: Optional[str] = None,
+            add_special_tokens: bool = False,
+        ) -> Dict[str, Any]:
+        encoded = self.tokenizer(
+            text,
+            max_length=max_length,
+            padding=padding,
+            truncation=truncation,
+            return_tensors=return_tensors,
+            add_special_tokens=add_special_tokens
+        )
+        return encoded
+
+    def encode_chat(
+            self,
+            text: List[Dict[str, Any]],
+        ) -> Dict[str, Any]:
+        encoded = self.tokenizer.apply_chat_template(
+            text,
+            tokenize=False
+        )
+        return encoded
+
+    def decode(
+        self,
+        token_ids: Union[List[int], List[List[int]]],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = False
+    ) -> Union[str, List[str]]:
+        if isinstance(token_ids[0], int):          # 单条
+            return self.tokenizer.decode(
+                token_ids,
+                skip_special_tokens=skip_special_tokens,
+                clean_up_tokenization_spaces=clean_up_tokenization_spaces
+            )
+        # batch
+        return self.tokenizer.batch_decode(
+            token_ids,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.tokenizer)
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.tokenizer.pad_token_id
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.tokenizer.eos_token_id
+
+    @property
+    def bos_token_id(self) -> int:
+        return self.tokenizer.bos_token_id
+
+    @property
+    def unk_token_id(self) -> int:
+        return self.tokenizer.unk_token_id
+
+    # 兼容旧接口
+    def get_vocab_size(self) -> int:
+        return self.vocab_size
+
+    def get_eot_token(self) -> int:
+        """End-Of-Text token，与 GPT 系列保持一致命名。"""
+        return self.eos_token_id
+
+if __name__ == '__main__':
+    tokenizer = Tokenizer("./llama2_tokenizer_32k")
+    print(tokenizer.get_eot_token())
+    encode = tokenizer.encode_chat([{"role": "system", "content": ""}, {"role": "user", "content": "请告诉我在中国古代的“四大发明”是什么？"}, {"role": "assistant", "content": "中国古代的“四大发明”是指造纸术、印刷术、火药和指南针。这四项发明对世界文明的发展产生了深远的影响：\n\n1. **造纸术**：据史书记载，东汉时期的蔡伦改进了造纸工艺，使得纸张的生产更加便捷、成本更低，质量也更加优良。这一发明极大地促进了文化的传播和保存。\n\n2. **印刷术**：中国古代的印刷术最早可以追溯到唐代的雕版印刷，到了宋代发展出了活字印刷技术。印刷术的发明极大地促进了知识的传播，降低了书籍的成本，对教育和文化的发展起到了重要的推动作用。\n\n3. **火药**：火药最初是在唐代被发现的，最初可能用于医疗或炼金术。到了宋代，人们开始将火药用于军事目的，发明了各种火器。火药的发明改变了战争的面貌，对世界军事技术的发展产生了重大影响。\n\n4. **指���针**：指南针最初在中国被用于风水测量，后来逐渐发展成为航海导航的重要工具。这一发明极大地促进了海上航行技术的发展，对于新航路的开辟和世界地理大发现起到了关键作用。\n\n这四项发明不仅在中国历史上占有重要地位，而且对全世界的科技进步和文明发展都产生了深远的影响。"}])
+    print(type(encode))
+    print("orin length", len(encode))
+    encode = tokenizer.encode(str(encode), max_length=None, truncation=True, padding="do_not_pad")
+    for token in encode['input_ids']:
+        print(tokenizer.decode([token]))
+    decode = tokenizer.decode([encode['input_ids']])
+    print("token lens: ", len(encode['input_ids']))
+    print("encode: ", encode)
+    print("att_mask: ", encode['attention_mask'])
+    print("decode: ", decode)
+    print("vocab_size", tokenizer.get_vocab_size())
+    print("eot", tokenizer.get_eot_token(), tokenizer.decode([tokenizer.get_eot_token()]))
+    print("eos", tokenizer.eos_token_id, tokenizer.decode([tokenizer.eos_token_id]))
+    print("unk", tokenizer.unk_token_id, tokenizer.decode([tokenizer.unk_token_id]))
+    print("pad", tokenizer.pad_token_id, tokenizer.decode([tokenizer.pad_token_id]))
+    print("bos", tokenizer.bos_token_id, tokenizer.decode([tokenizer.bos_token_id]))
+    for i in range(10):
+        print(i, tokenizer.decode([i]))
+
+
+
+
+
+
+
+
+
+