qwen2:infer初步实现推理,待分析问题

Author: miaobyte

model/safetensor_deepx/safetensor_deepx/loader.py

@@ -5,6 +5,8 @@
 import yaml
 import argparse
 import shutil
+import glob
+import re
 
 
 class TensorInfo:
@@ -54,19 +56,37 @@ def _load_config(self):
                 return json.load(f)
         return {}
 
+    def _find_model_files(self):
+        """查找所有分片模型文件"""
+        single_file = os.path.join(self.model_dir, "model.safetensors")
+        shard_files = glob.glob(os.path.join(self.model_dir, "model-*-of-*.safetensors"))
+        
+        # 使用正则表达式提取分片编号
+        pattern = re.compile(r"model-(\d+)-of-(\d+)\.safetensors")
+        filtered_shards = []
+        for f in shard_files:
+            match = pattern.search(os.path.basename(f))
+            if match:
+                filtered_shards.append( (int(match.group(1)), f) )
+        
+        if os.path.exists(single_file):
+            return [single_file]
+        elif filtered_shards:
+            # 按分片编号排序后返回路径
+            filtered_shards.sort(key=lambda x: x[0])
+            return [f[1] for f in filtered_shards]
+        raise FileNotFoundError(f"No model files found in {self.model_dir}")
+
     def export(self):
         """导出safetensor模型到指定目录"""
-        model_path = os.path.join(self.model_dir, "model.safetensors")
-        if not os.path.exists(model_path):
-            raise FileNotFoundError(f"找不到模型文件: {model_path}")
-
-        # 修改为使用PyTorch框架加载
-        with safe_open(model_path, framework="pt") as f:  # 改为pt框架
-            for key in f.keys():
-                tensor = f.get_tensor(key)
-                self._save_tensor(key, tensor)
+        model_files = self._find_model_files()
+        
+        for model_path in model_files:
+            with safe_open(model_path, framework="pt") as f:
+                for key in f.keys():
+                    tensor = f.get_tensor(key)
+                    self._save_tensor(key, tensor)
 
-        # 保存全局配置
         self._save_config()
         self._copy_tokenizer_files()
 
@@ -135,34 +155,57 @@ def _load_config(self):
                 return json.load(f)
         return {}
 
+    def _find_model_files(self):
+        """查找所有分片模型文件"""
+        single_file = os.path.join(self.model_dir, "model.safetensors")
+        shard_files = glob.glob(os.path.join(self.model_dir, "model-*-of-*.safetensors"))
+        
+        # 统一使用正则表达式匹配
+        pattern = re.compile(r"model-(\d+)-of-(\d+)\.safetensors")
+        filtered_shards = []
+        for f in shard_files:
+            match = pattern.search(os.path.basename(f))
+            if match:
+                filtered_shards.append( (int(match.group(1)), f) )
+        
+        if os.path.exists(single_file):
+            return [single_file]
+        elif filtered_shards:
+            filtered_shards.sort(key=lambda x: x[0])
+            return [f[1] for f in filtered_shards]
+        else:
+            raise FileNotFoundError(f"No model files found in {self.model_dir}")
+
     def load(self):
         """加载safetensor模型文件"""
         tensors = {}
         metadata = {}
-
-        model_path = os.path.join(self.model_dir, "model.safetensors")
-        if not os.path.exists(model_path):
-            raise FileNotFoundError(f"找不到模型文件: {model_path}")
-
-        with safe_open(model_path, framework="pt") as f:  # 修改为pt框架
-            metadata = f.metadata() if hasattr(f, 'metadata') else {}
-            for key in f.keys():
-                pt_tensor = f.get_tensor(key).cpu().detach()  # 获取PyTorch张量
-
-                # 构造TensorInfo
-                tensor_info = TensorInfo(
-                    dtype=str(pt_tensor.dtype).replace("torch.", ""),
-                    ndim=pt_tensor.ndim,
-                    shape=tuple(pt_tensor.shape),
-                    size=pt_tensor.numel(),
-                    strides=pt_tensor.stride() if pt_tensor.is_contiguous() else None
-                )
-
-                # 转换为字节流（保持内存对齐）
-                byte_buffer = pt_tensor.numpy().tobytes() if pt_tensor.device == "cpu" \
-                    else pt_tensor.cpu().numpy().tobytes()
-
-                tensors[key] = Tensor(byte_buffer, tensor_info)
+        
+        model_files = self._find_model_files()
+        
+        for model_path in model_files:
+            with safe_open(model_path, framework="pt") as f:
+                # 合并metadata
+                file_metadata = f.metadata() if hasattr(f, 'metadata') else {}
+                metadata.update(file_metadata)
+                
+                for key in f.keys():
+                    pt_tensor = f.get_tensor(key).cpu().detach()
+
+                    # 构造TensorInfo
+                    tensor_info = TensorInfo(
+                        dtype=str(pt_tensor.dtype).replace("torch.", ""),
+                        ndim=pt_tensor.ndim,
+                        shape=tuple(pt_tensor.shape),
+                        size=pt_tensor.numel(),
+                        strides=pt_tensor.stride() if pt_tensor.is_contiguous() else None
+                    )
+
+                    # 转换为字节流（保持内存对齐）
+                    byte_buffer = pt_tensor.numpy().tobytes() if pt_tensor.device == "cpu" \
+                        else pt_tensor.cpu().numpy().tobytes()
+
+                    tensors[key] = Tensor(byte_buffer, tensor_info)
 
         metadata["model_config"] = self.config
         return tensors, metadata

todo/infer.py

@@ -0,0 +1,129 @@
+import sys
+import threading
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+import torch
+
+def init_model():
+    model_path = "/home/lipeng/model/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        tokenizer.pad_token = tokenizer.eos_token
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            trust_remote_code=True,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            # use_flash_attention_2=True  # 启用Flash Attention
+        ).eval()
+        
+        return model, tokenizer
+    except Exception as e:
+        raise RuntimeError(f"模型初始化失败: {str(e)}")
+
+class StdoutStreamer(TextStreamer):
+    def __init__(self, tokenizer):
+        super().__init__(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        self.cache = []
+        self.first_token = True
+    
+    def on_finalized_text(self, text: str, stream_end: bool = False):
+        self.cache.append(text)
+        if stream_end or len(self.cache) >= 2:
+            full_text = "".join(self.cache)
+            sys.stdout.write(full_text)
+            sys.stdout.flush()
+            self.cache = []
+
+def generate_stream(model, tokenizer, text, max_length):
+    formatted_text = f"<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
+    inputs = tokenizer(
+        formatted_text, 
+        return_tensors='pt', 
+        add_special_tokens=False,
+        return_attention_mask=True
+    ).to(model.device)
+    streamer = StdoutStreamer(tokenizer)
+    
+    generation_kwargs = {
+        "input_ids": inputs.input_ids,
+        "attention_mask": inputs.attention_mask,
+        "max_new_tokens": max_length,
+        "pad_token_id": tokenizer.eos_token_id,
+        "temperature": 0.3,      # 降低随机性
+        "top_p": 0.85,           # 限制采样范围
+        "repetition_penalty": 1.2, # 增强重复抑制
+        "do_sample": True,
+        "streamer": streamer
+    }
+    
+    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    thread.join()
+    print("\n")  # 流式结束换行
+
+def generate_text(model, tokenizer, text, max_length=50):
+    formatted_text = f"<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
+    inputs = tokenizer(
+        formatted_text,
+        return_tensors='pt',
+        add_special_tokens=False,
+        return_attention_mask=True
+    ).to(model.device)
+    
+    with torch.no_grad():
+        output = model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=max_length,
+            pad_token_id=tokenizer.eos_token_id,
+            temperature=0.3,
+            top_p=0.85,
+            repetition_penalty=1.2,
+            do_sample=True
+        )
+    
+    return tokenizer.decode(
+        output[0][len(inputs.input_ids[0]):],
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True
+    )
+
+def main():
+    try:
+        model, tokenizer = init_model()
+        sys.stderr.write("模型加载成功，输入提示开始生成（Ctrl+C退出）\n")
+    except Exception as e:
+        sys.stderr.write(f"服务启动失败: {e}\n")
+        return
+
+    # 单独测试分词器
+    text = "<|im_start|>user\n你好<|im_end|>\n<|im_start|>assistant\n"
+    tokens = tokenizer.encode(text, add_special_tokens=False)
+    decoded = tokenizer.decode(tokens)
+    assert decoded == text  # 验证编码解码一致性
+    try:
+        for line in sys.stdin:
+            text = line.strip()
+            if not text:
+                continue
+
+            # 固定参数设置
+            max_length = 2048  # 最大生成长度
+            stream = True     # 始终使用流式
+            
+            if stream:
+                generate_stream(model, tokenizer, text, max_length)
+            else:
+                result = generate_text(model, tokenizer, text, max_length)
+                print(result)
+                
+    except KeyboardInterrupt:
+        sys.stderr.write("\n服务已终止\n")
+    except Exception as e:
+        sys.stderr.write(f"运行时错误: {str(e)}\n")
+
+if __name__ == '__main__':
+    main()
+
+

todo/qwen2_infer.py

@@ -8,7 +8,6 @@
 import torch
 import torch.nn as nn
 from transformers import AutoTokenizer
-import re
 
 class ModelConfig:
     def __init__(self):