Update infer.py

infer.py

@@ -1,7 +1,3 @@
-"""
-Flask推理界面 - 多模态Dense Transformer (适配 Qwen Tokenizer 版)
-"""
-
 import os
 import torch
 import torch.nn.functional as F
@@ -14,18 +10,12 @@ import base64
 from pathlib import Path
 from typing import Optional
 
-# 确保引入路径正确，根据你之前的文件结构
 from model import MultiModalDenseTransformer
-# 注意：UnifiedMultiModalPreprocessor 之前是在 continual_learning.py 中定义的
-# 如果你移动了它，请修改这里的导入路径
 from continual_learning import UnifiedMultiModalPreprocessor
-# 如果没有 image_transform，我们需要在这里定义或导入
 from torchvision import transforms
 
-# 设置国内镜像
 os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
 
-# 定义图像预处理 (与 training 保持一致)
 image_transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -33,8 +23,6 @@ image_transform = transforms.Compose([
 ])
 
 class ModelInference:
-    """模型推理类"""
-    
     def __init__(
         self,
         checkpoint_path: str,
@@ -44,8 +32,6 @@ class ModelInference:
     ):
         self.device = torch.device(device)
         print(f"Using device: {self.device}")
-        
-        # 1. 加载 Tokenizer (与预训练一致)
         print(f"Loading tokenizer: {tokenizer_name}...")
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(
@@ -59,26 +45,24 @@ class ModelInference:
         except Exception as e:
             print(f"Error loading tokenizer: {e}")
             raise e
-        
-        # 2. 配置模型参数 (必须与 pretrain.py 中的配置完全一致)
+
         if config_path and Path(config_path).exists():
             with open(config_path, 'r') as f:
                 self.config = json.load(f)
         else:
-            # [CRITICAL] 这里使用了你在 pretrain.py 中使用的参数
             self.config = {
-                'model_dim': 1536,           # 预训练设置
-                'vocab_size': len(self.tokenizer), # 自动适配 Qwen (约 151665)
-                'n_layers': 12,             # 预训练设置
-                'n_heads': 12,              # 预训练设置
-                'n_kv_heads': 4,            # 预训练设置
-                'head_dim': None,           # 自动计算
-                'max_seq_len': 512,        # 预训练设置
-                'dropout': 0.0,             # 推理时关闭 dropout
-                'use_moe': False,           # 预训练设置
-                'use_adapter': False,       # 预训练未开启 Adapter
-                'use_lora': False,          # 预训练未开启 LoRA
-                'rope_scaling_type': "yarn" # 预训练设置
+                'model_dim': 1536,           
+                'vocab_size': len(self.tokenizer), 
+                'n_layers': 12,
+                'n_heads': 12,         
+                'n_kv_heads': 4,      
+                'head_dim': None,          
+                'max_seq_len': 512,       
+                'dropout': 0.0,             
+                'use_moe': False,          
+                'use_adapter': False,       
+                'use_lora': False,          
+                'rope_scaling_type': "yarn" 
             }
         
         # 3. 初始化模型结构
@@ -91,21 +75,18 @@ class ModelInference:
             
             # 4. 加载权重
             print(f"Loading checkpoint from {checkpoint_path}...")
-            # weights_only=False 是为了支持加载完整的 checkpoint 字典
             checkpoint = torch.load(
                 checkpoint_path, 
                 map_location=self.device,
                 weights_only=False 
             )
             
-            # 提取 state_dict
             if 'model_state_dict' in checkpoint:
                 print("Found 'model_state_dict' in checkpoint.")
                 state_dict = checkpoint['model_state_dict']
             else:
                 state_dict = checkpoint
             
-            # 处理可能的键名不匹配 (如 DDP 训练产生的 'module.' 前缀)
             new_state_dict = {}
             for k, v in state_dict.items():
                 if k.startswith('module.'):
@@ -113,7 +94,6 @@ class ModelInference:
                 else:
                     new_state_dict[k] = v
             
-            # 加载权重 (strict=False 允许忽略一些非关键的不匹配，如 loss 缓存等)
             missing, unexpected = self.model.load_state_dict(new_state_dict, strict=False)
             if missing:
                 print(f"Warning: Missing keys: {len(missing)}")
@@ -143,40 +123,29 @@ class ModelInference:
         image: Optional[Image.Image] = None
     ) -> str:
         """生成文本"""
-        
-        # 编码输入
         inputs = self.tokenizer(prompt, return_tensors="pt")
         input_ids = inputs['input_ids'].to(self.device)
-        
-        # 构建 MultiModalDenseTransformer 需要的输入格式
         input_data = {'segments': []}
         
         # 处理图像
         if image is not None:
             if image.mode != 'RGB':
                 image = image.convert('RGB')
-            # 简单的图像处理
             image_tensor = image_transform(image).unsqueeze(0).to(self.device)
-            # 这里假设预处理器能处理这种输入
             try:
-                # process_batch 接受 (batch_data, modality_type) 并返回 segments 列表
                 mod_segments = self.preprocessor.process_batch(image_tensor, 'image')
-                # 将返回的 segment 列表合并到 input_data
                 for seg in mod_segments:
                     input_data['segments'].append(seg)
             except Exception as e:
                 print(f"Warning: Image processing skipped due to error: {e}")
-        
-        # 添加文本段
+    
         input_data['segments'].append({
             'type': 'text',
             'data': input_ids,
             'modality_id': 0
         })
         
-        # 生成
         try:
-            # 使用模型自带的 generate 方法
             generated_ids = self.model.generate(
                 input_data,
                 max_new_tokens=max_new_tokens,
@@ -188,19 +157,11 @@ class ModelInference:
                 eos_token_id=self.tokenizer.eos_token_id,
                 pad_token_id=self.tokenizer.pad_token_id
             )
-            
-            # 解码
-            # 注意：生成的 ids 可能包含原始输入，或者只包含新生成的 token
-            # MultiModalDenseTransformer.generate 通常返回完整的序列
+
             generated_text = self.tokenizer.decode(
                 generated_ids[0],
                 skip_special_tokens=True
             )
-            
-            # 如果包含 prompt，可以选择移除它只显示新内容
-            # if generated_text.startswith(prompt):
-            #     generated_text = generated_text[len(prompt):]
-            
             return generated_text
             
         except Exception as e:
@@ -209,8 +170,6 @@ class ModelInference:
             traceback.print_exc()
             return f"Error: {str(e)}"
 
-
-# 全局模型实例
 model_instance = None
 app = Flask(__name__)
 app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
@@ -274,7 +233,7 @@ def create_html_template():
 </head>
 <body>
     <div class="container">
-        <h1>🚀 模型在线推理</h1>
+        <h1> 模型在线推理</h1>
         
         <div>
             <label><strong>提示词 (Prompt):</strong></label>