Update infer.py

infer.py

@@ -1,150 +1,176 @@
 import os
-import torch
-import torch.nn.functional as F
-from flask import Flask, render_template, request, jsonify
-from transformers import AutoTokenizer
-from PIL import Image
-import json
-import io
-import base64
+import argparse
 from pathlib import Path
+import json
 from typing import Optional
 
+import torch
+from PIL import Image
+from transformers import AutoTokenizer
+
+import gradio as gr
 from model import MultiModalDenseTransformer
 from continual_learning import UnifiedMultiModalPreprocessor
-from torchvision import transforms
-
 os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
 
+from torchvision import transforms
 image_transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225]),
 ])
-
 class ModelInference:
     def __init__(
-        self,
-        checkpoint_path: str,
-        tokenizer_name: str,
-        config_path: Optional[str] = None,
+        self, 
+        checkpoint_path: str, 
+        tokenizer_name: str, 
+        config_path: Optional[str] = None, 
         device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
     ):
         self.device = torch.device(device)
-        print(f"Using device: {self.device}")
-        print(f"Loading tokenizer: {tokenizer_name}...")
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                tokenizer_name, 
-                use_fast=True, 
-                trust_remote_code=True
-            )
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
-        except Exception as e:
-            print(f"Error loading tokenizer: {e}")
-            raise e
-
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_name, 
+            use_fast=True, 
+            trust_remote_code=True
+        )
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
         if config_path and Path(config_path).exists():
             with open(config_path, 'r') as f:
                 self.config = json.load(f)
         else:
             self.config = {
-                'model_dim': 1536,           
-                'vocab_size': len(self.tokenizer), 
+                '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" 
+                '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",
+                'use_multimodal_fusion': False,
+                'use_contrastive': False
             }
+
+        self.model = MultiModalDenseTransformer(**self.config)
+        self.preprocessor = UnifiedMultiModalPreprocessor(model_dim=self.config['model_dim'])
+
+        print(f"Loading checkpoint from {checkpoint_path}...")
+        checkpoint = torch.load(checkpoint_path, map_location=self.device)
         
-        # 3. 初始化模型结构
-        print("Initializing model architecture...")
-        try:
-            self.model = MultiModalDenseTransformer(**self.config)
-            self.preprocessor = UnifiedMultiModalPreprocessor(
-                model_dim=self.config['model_dim']
-            )
-            
-            # 4. 加载权重
-            print(f"Loading checkpoint from {checkpoint_path}...")
-            checkpoint = torch.load(
-                checkpoint_path, 
-                map_location=self.device,
-                weights_only=False 
-            )
-            
-            if 'model_state_dict' in checkpoint:
-                print("Found 'model_state_dict' in checkpoint.")
-                state_dict = checkpoint['model_state_dict']
+        state_dict = None
+        
+        if 'actor_state_dict' in checkpoint:
+            print("Detected GRPO checkpoint format (actor_state_dict)")
+            state_dict = checkpoint['actor_state_dict']
+        elif 'model_state_dict' in checkpoint:
+            print("Detected Standard/SFT checkpoint format (model_state_dict)")
+            state_dict = checkpoint['model_state_dict']
+        else:
+            print("Detected raw state dict format")
+            state_dict = checkpoint
+
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if k.startswith('module.'):
+                new_state_dict[k[7:]] = v
             else:
-                state_dict = checkpoint
-            
-            new_state_dict = {}
-            for k, v in state_dict.items():
-                if k.startswith('module.'):
-                    new_state_dict[k[7:]] = v
-                else:
-                    new_state_dict[k] = v
-            
-            missing, unexpected = self.model.load_state_dict(new_state_dict, strict=False)
-            if missing:
-                print(f"Warning: Missing keys: {len(missing)}")
-            if unexpected:
-                print(f"Warning: Unexpected keys: {len(unexpected)}")
-            
-            self.model.to(self.device)
-            self.preprocessor.to(self.device)
-            self.model.eval()
-            
-            print("Model loaded successfully!")
-            print(f"Total parameters: {sum(p.numel() for p in self.model.parameters()) / 1e6:.2f}M")
-            
-        except Exception as e:
-            print(f"Error initializing model: {e}")
-            raise e
-    
+                new_state_dict[k] = v
+
+        missing, unexpected = self.model.load_state_dict(new_state_dict, strict=False)
+        if missing:
+            print(f"Warning: Missing keys: {len(missing)}")
+            if len(missing) <= 10:
+                print(f"Missing keys: {missing}")
+        if unexpected:
+            print(f"Warning: Unexpected keys: {len(unexpected)}")
+            if len(unexpected) <= 10:
+                print(f"Unexpected keys: {unexpected}")
+
+        self.model.to(self.device)
+        self.preprocessor.to(self.device)
+        self.model.eval()
+
+
+    def _build_position_ids(self, attention_mask: torch.Tensor) -> torch.Tensor:
+
+        batch_size, seq_len = attention_mask.shape
+        position_ids = torch.zeros((batch_size, seq_len), dtype=torch.long, device=self.device)
+        
+        for i in range(batch_size):
+            non_pad_positions = (attention_mask[i] == 1).nonzero(as_tuple=True)[0]
+            if len(non_pad_positions) > 0:
+                start_pos = non_pad_positions[0].item()
+                valid_len = len(non_pad_positions)
+                # 从 0 开始编号有效 token 的位置
+                position_ids[i, start_pos:start_pos + valid_len] = torch.arange(
+                    valid_len, 
+                    device=self.device
+                )
+        
+        return position_ids
+
     @torch.no_grad()
     def generate_text(
-        self,
-        prompt: str,
-        max_new_tokens: int = 128,
-        temperature: float = 0.7,
-        top_k: int = 40,
-        top_p: float = 0.9,
-        repetition_penalty: float = 1.1,
+        self, 
+        prompt: str, 
+        max_new_tokens: int = 128, 
+        temperature: float = 0.7, 
+        top_k: int = 40, 
+        top_p: float = 0.9, 
+        repetition_penalty: float = 1.1, 
         image: Optional[Image.Image] = None
     ) -> str:
-        """生成文本"""
-        inputs = self.tokenizer(prompt, return_tensors="pt")
-        input_ids = inputs['input_ids'].to(self.device)
-        input_data = {'segments': []}
+        formatted_prompt = f"user: {prompt}\nassistant:\n<think>\n"
         
-        # 处理图像
+        inputs = self.tokenizer(
+            formatted_prompt, 
+            return_tensors="pt",
+            padding=False  
+        )
+        input_ids = inputs['input_ids'].to(self.device)
+        attention_mask = inputs['attention_mask'].to(self.device)
+
+        segments = []
+
+        segments.append({
+            'type': 'text', 
+            'data': input_ids, 
+            'modality_id': 0
+        })
+
+        has_image = False
         if image is not None:
-            if image.mode != 'RGB':
-                image = image.convert('RGB')
-            image_tensor = image_transform(image).unsqueeze(0).to(self.device)
             try:
-                mod_segments = self.preprocessor.process_batch(image_tensor, 'image')
-                for seg in mod_segments:
-                    input_data['segments'].append(seg)
+                if image.mode != 'RGB':
+                    image = image.convert('RGB')
+                image_tensor = image_transform(image).unsqueeze(0).to(self.device)
+                
+                segments.append({
+                    'type': 'image',
+                    'data': image_tensor,
+                    'modality_id': 1
+                })
+                has_image = True
+                print("Image added to input")
+                
             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
-        })
-        
+                print(f"Warning: Image processing error: {e}")
+
+        position_ids = self._build_position_ids(attention_mask)
+        input_data = {
+            'segments': segments,
+        }
+        input_data['attention_mask'] = attention_mask
+        if not has_image:
+            input_data['position_ids'] = position_ids
+
         try:
             generated_ids = self.model.generate(
                 input_data,
@@ -158,174 +184,191 @@ class ModelInference:
                 pad_token_id=self.tokenizer.pad_token_id
             )
 
-            generated_text = self.tokenizer.decode(
-                generated_ids[0],
-                skip_special_tokens=True
-            )
-            return generated_text
+            output_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+            return output_text.strip()
             
         except Exception as e:
-            print(f"Generation error: {e}")
             import traceback
             traceback.print_exc()
-            return f"Error: {str(e)}"
-
-model_instance = None
-app = Flask(__name__)
-app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
-
-@app.route('/')
-def index():
-    display_config = model_instance.config.copy() if model_instance else {}
-    return render_template('index.html', config=display_config)
+            return f"Error during generation: {str(e)}"
 
-@app.route('/generate', methods=['POST'])
-def generate():
-    try:
-        data = request.json
-        prompt = data.get('prompt', '')
-        if not prompt.strip():
-            return jsonify({'error': '请输入提示文本'}), 400
+def build_ui(model_instance):
+    with gr.Blocks(title="MultiModal Dense Transformer - Gradio", css="""
+        .gradio-container { max-width: 900px; margin: auto; }
+    """) as demo:
+        gr.Markdown("##  在线推理（文本）")
         
-        max_tokens = int(data.get('max_tokens', 100))
-        temperature = float(data.get('temperature', 0.7))
-        top_k = int(data.get('top_k', 40))
-        top_p = float(data.get('top_p', 0.9))
-        repetition_penalty = float(data.get('repetition_penalty', 1.1))
-        
-        image = None
-        if 'image' in data and data['image']:
-            try:
-                image_data = base64.b64decode(data['image'].split(',')[1])
-                image = Image.open(io.BytesIO(image_data))
-            except Exception as e:
-                print(f"Image load error: {e}")
-        
-        output = model_instance.generate_text(
-            prompt, max_tokens, temperature, top_k, top_p, repetition_penalty, image
+        with gr.Row():
+            with gr.Column(scale=3):
+                txt = gr.Textbox(
+                    label="Prompt (Instruction)", 
+                    placeholder="请输入指令或问题...", 
+                    lines=5
+                )
+                img = gr.Image(type="pil", label="(可选) 上传图片（支持多模态）")
+                btn = gr.Button("生成 (Generate)", variant="primary")
+                
+            with gr.Column(scale=2):
+                max_tokens = gr.Slider(
+                    label="Max New Tokens", 
+                    minimum=16, 
+                    maximum=1024, 
+                    step=1, 
+                    value=128
+                )
+                temperature = gr.Slider(
+                    label="Temperature", 
+                    minimum=0.1, 
+                    maximum=1.5, 
+                    step=0.01, 
+                    value=0.7
+                )
+                top_k = gr.Slider(
+                    label="Top-k", 
+                    minimum=0, 
+                    maximum=200, 
+                    step=1, 
+                    value=40
+                )
+                top_p = gr.Slider(
+                    label="Top-p", 
+                    minimum=0.0, 
+                    maximum=1.0, 
+                    step=0.01, 
+                    value=0.9
+                )
+                rep_pen = gr.Slider(
+                    label="Repetition Penalty", 
+                    minimum=0.5, 
+                    maximum=2.0, 
+                    step=0.01, 
+                    value=1.1
+                )
+                status = gr.Textbox(
+                    label="Status", 
+                    value="Ready", 
+                    interactive=False
+                )
+                
+        output = gr.Textbox(label="Output", lines=12, interactive=False)
+        gr.Examples(
+            examples=[
+                ["请解释什么是深度学习", None],
+                ["计算 123 + 456 等于多少？", None],
+                ["写一首关于春天的诗", None],
+            ],
+            inputs=[txt, img],
         )
-        return jsonify({'output': output})
-    
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
 
-def create_html_template():
-    """写入HTML模板"""
-    html_content = '''
-<!DOCTYPE html>
-<html lang="zh-CN">
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>Model Inference</title>
-    <style>
-        body { font-family: sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; background: #f0f2f5; }
-        .container { background: white; padding: 30px; border-radius: 12px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
-        h1 { color: #1a73e8; text-align: center; }
-        textarea { width: 100%; padding: 10px; border: 1px solid #ddd; border-radius: 8px; margin: 10px 0; min-height: 100px; }
-        .controls { display: grid; grid-template-columns: 1fr 1fr; gap: 15px; margin: 20px 0; background: #f8f9fa; padding: 15px; border-radius: 8px; }
-        button { background: #1a73e8; color: white; border: none; padding: 12px 24px; border-radius: 6px; cursor: pointer; width: 100%; font-size: 16px; transition: background 0.3s; }
-        button:hover { background: #1557b0; }
-        button:disabled { background: #ccc; }
-        #output { margin-top: 20px; padding: 20px; background: #f8f9fa; border-radius: 8px; white-space: pre-wrap; min-height: 100px; border: 1px solid #e0e0e0; }
-        .loading { color: #666; font-style: italic; }
-    </style>
-</head>
-<body>
-    <div class="container">
-        <h1> 模型在线推理</h1>
-        
-        <div>
-            <label><strong>提示词 (Prompt):</strong></label>
-            <textarea id="prompt" placeholder="请输入你的问题..."></textarea>
-        </div>
+        def gr_generate(prompt, image, max_tokens_v, temp_v, topk_v, topp_v, rep_v):
+            if not prompt or str(prompt).strip() == "":
+                return "", " 请输入 Prompt"
+            
+            try:
+                status_msg = " Generating..."
+                # 调用模型生成
+                out = model_instance.generate_text(
+                    prompt=prompt,
+                    max_new_tokens=int(max_tokens_v),
+                    temperature=float(temp_v),
+                    top_k=int(topk_v),
+                    top_p=float(topp_v),
+                    repetition_penalty=float(rep_v),
+                    image=image
+                )
+                return out, " Done"
+            except Exception as e:
+                return f"Error: {str(e)}", " Error"
 
-        <div class="controls">
-            <div>
-                <label>Max Tokens: <span id="maxTokensVal">128</span></label>
-                <input type="range" id="maxTokens" min="32" max="1024" value="128" style="width:100%" oninput="document.getElementById('maxTokensVal').innerText=this.value">
-            </div>
-            <div>
-                <label>Temperature: <span id="tempVal">0.7</span></label>
-                <input type="range" id="temperature" min="0.1" max="1.5" step="0.1" value="0.7" style="width:100%" oninput="document.getElementById('tempVal').innerText=this.value">
-            </div>
-        </div>
+        btn.click(
+            fn=gr_generate, 
+            inputs=[txt, img, max_tokens, temperature, top_k, top_p, rep_pen], 
+            outputs=[output, status]
+        )
 
-        <button id="btn" onclick="generate()">生成 (Generate)</button>
-        
-        <div id="output">结果将显示在这里...</div>
-    </div>
-
-    <script>
-        async function generate() {
-            const prompt = document.getElementById('prompt').value;
-            if(!prompt) return alert("请输入内容");
-            
-            const btn = document.getElementById('btn');
-            const out = document.getElementById('output');
-            
-            btn.disabled = true;
-            btn.innerText = "生成中...";
-            out.innerHTML = '<div class="loading">正在思考中...</div>';
-            
-            try {
-                const res = await fetch('/generate', {
-                    method: 'POST',
-                    headers: {'Content-Type': 'application/json'},
-                    body: JSON.stringify({
-                        prompt: prompt,
-                        max_tokens: parseInt(document.getElementById('maxTokens').value),
-                        temperature: parseFloat(document.getElementById('temperature').value)
-                    })
-                });
-                const data = await res.json();
-                if(data.error) out.innerText = "Error: " + data.error;
-                else out.innerText = data.output;
-            } catch(e) {
-                out.innerText = "请求失败: " + e;
-            } finally {
-                btn.disabled = false;
-                btn.innerText = "生成 (Generate)";
-            }
-        }
-    </script>
-</body>
-</html>
-    '''
-    
-    Path('templates').mkdir(exist_ok=True)
-    with open('templates/index.html', 'w', encoding='utf-8') as f:
-        f.write(html_content)
+    return demo
 
 def main():
-    import argparse
-    parser = argparse.ArgumentParser()
-    # 默认指向 pretrain 保存的 checkpoint 路径
-    parser.add_argument("--checkpoint", type=str, default="/root/multimodal/checkpoints/pretrain_fixed/step_10000.pt") 
-    parser.add_argument("--tokenizer", type=str, default="Qwen/Qwen2.5-7B-Instruct")
-    parser.add_argument("--port", type=int, default=5001)
-    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser = argparse.ArgumentParser(
+        description="Gradio inference interface for MultiModal Dense Transformer"
+    )
+    parser.add_argument(
+        "--checkpoint", 
+        type=str, 
+        default="/root/checkpoints/dcpo_posttrain_round3/step_15600.pt",
+        help="Path to model checkpoint"
+    )
+    parser.add_argument(
+        "--tokenizer", 
+        type=str, 
+        default="Qwen/Qwen2.5-7B-Instruct",
+        help="Tokenizer name or path"
+    )
+    parser.add_argument(
+        "--config", 
+        type=str, 
+        default=None,
+        help="Path to model config JSON (optional)"
+    )
+    parser.add_argument(
+        "--port", 
+        type=int, 
+        default=5001,
+        help="Port to run Gradio server"
+    )
+    parser.add_argument(
+        "--share", 
+        type=lambda x: x.lower() in ("true","1","yes"), 
+        default=True,
+        help="Create public link (True/False)"
+    )
     args = parser.parse_args()
-    
+
     if not Path(args.checkpoint).exists():
-        # 尝试找最近的 step checkpoint
-        steps = list(Path("checkpoints/pretrain").glob("step_*.pt"))
-        if steps:
-            print(f"未找到 final_model.pt，尝试使用最新的 checkpoint: {steps[-1]}")
-            args.checkpoint = str(steps[-1])
-        else:
-            print(f"错误: 找不到检查点文件: {args.checkpoint}")
-            return
+        print(f" Checkpoint not found: {args.checkpoint}")
+        
+        possible_dirs = [
+            Path("/root/checkpoints/posttrain/grpo"),
+            Path("/root/checkpoints/dcpo_training"),
+            Path("/root/checkpoints/r1_zero_reproduction"),
+        ]
+        
+        for checkpoint_dir in possible_dirs:
+            if checkpoint_dir.exists():
+                grpo_files = sorted(
+                    [p for p in checkpoint_dir.glob("grpo_iter_*.pt")], 
+                    key=lambda p: int(p.stem.split('_')[-1]) if p.stem.split('_')[-1].isdigit() else 0
+                )
+
+                step_files = sorted(
+                    [p for p in checkpoint_dir.glob("step_*.pt")], 
+                    key=lambda p: int(p.stem.split('_')[-1]) if p.stem.split('_')[-1].isdigit() else 0
+                )
+                
+                candidates = grpo_files + step_files
+                if candidates:
+                    args.checkpoint = str(candidates[-1])
+                    print(f" Using latest checkpoint: {args.checkpoint}")
+                    break
+        
+        if not Path(args.checkpoint).exists():
+            raise FileNotFoundError(f"找不到可用的检查点文件")
 
-    create_html_template()
     
     global model_instance
-    model_instance = ModelInference(args.checkpoint, args.tokenizer)
+    model_instance = ModelInference(
+        args.checkpoint, 
+        args.tokenizer, 
+        args.config
+    )
+
     
-    print(f"\n服务已启动: http://{args.host}:{args.port}")
-    app.run(host=args.host, port=args.port,
-    debug=True,  # 开启调试模式
-    use_reloader=False)
+    demo = build_ui(model_instance)
+    demo.launch(
+        server_port=args.port, 
+        share=args.share,
+        server_name="0.0.0.0"  # 允许外部访问
+    )
 
 if __name__ == "__main__":
     main()