#!/opt/homebrew/bin/python3 """ PaddleOCR-VL MLX 优化版 - 加载视觉编码器权重 提升准确度 作者: AI Assistant 日期: 2024-12-25 版本: v6.0 - 优化准确度 """ import mlx.core as mx import mlx.nn as nn from PIL import Image, ImageDraw import numpy as np import json from pathlib import Path from typing import Optional, List, Tuple import time # 导入基础组件 from mlx_components import ( RMSNorm, Attention, MLP, DecoderLayer ) class ImprovedPaddleOCRMLX: """改进的 PaddleOCR MLX - 加载视觉编码器权重""" def __init__(self, model_dir: str): self.model_dir = Path(model_dir) print("🚀 初始化改进版 PaddleOCR MLX...") print(f"📂 模型目录: {model_dir}") # 加载配置 self.config = self._load_config() # 加载 tokenizer self.tokenizer = self._load_tokenizer() # 创建模型 self.model = self._create_model() # 加载所有权重（包括视觉编码器） self._load_all_weights() print("✅ 初始化完成！") def _load_config(self) -> dict: """加载模型配置""" config_path = self.model_dir / "config.json" with open(config_path, 'r') as f: config = json.load(f) print(f"✅ 配置加载完成") return config def _load_tokenizer(self): """加载 tokenizer""" try: from transformers import AutoTokenizer original_model_path = "/Users/gt/.lmstudio/hub/models/paddleocr-vl" tokenizer = AutoTokenizer.from_pretrained( original_model_path, trust_remote_code=True ) print(f"✅ Tokenizer 加载完成 (词汇表: {len(tokenizer)})") return tokenizer except Exception as e: print(f"⚠️ Tokenizer 加载失败: {e}") return None def _create_model(self): """创建模型""" print("🔄 创建模型...") class ImprovedModel(nn.Module): """改进的模型 - 包含完整的视觉编码器""" def __init__(self, config): super().__init__() self.config = config # 语言模型配置 self.hidden_size = config.get('hidden_size', 1024) self.vocab_size = config.get('vocab_size', 103424) self.intermediate_size = config.get('intermediate_size', 3072) self.num_attention_heads = config.get('num_attention_heads', 16) self.num_kv_heads = config.get('num_key_value_heads', 2) self.num_hidden_layers = config.get('num_hidden_layers', 18) self.head_dim = config.get('head_dim', 128) # 视觉编码器配置 vision_config = config.get('vision_config', {}) self.vision_hidden_size = vision_config.get('hidden_size', 1152) self.vision_num_layers = 27 # 从权重中看到的 # 视觉编码器 # Patch embedding (卷积) self.patch_embedding = nn.Conv2d( in_channels=3, out_channels=self.vision_hidden_size, kernel_size=14, stride=14, bias=True ) # 位置嵌入 self.position_embedding = mx.zeros((729, self.vision_hidden_size)) # 视觉 Transformer 层 self.vision_layers = [ DecoderLayer( hidden_size=self.vision_hidden_size, num_heads=16, intermediate_size=4304, num_kv_heads=16, # 视觉编码器使用标准注意力 head_dim=72, # 1152 / 16 = 72 ) for _ in range(self.vision_num_layers) ] # 视觉编码器的归一化 self.vision_norm = RMSNorm(self.vision_hidden_size) # 视觉投影层 (mlp_AR) # spatial_merge (2x2) + pre_norm + linear_1 + linear_2 self.vision_pre_norm = nn.LayerNorm(self.vision_hidden_size) self.vision_linear_1 = nn.Linear(4608, 4608, bias=True) self.vision_linear_2 = nn.Linear(4608, self.hidden_size, bias=True) # Token 嵌入 self.embed_tokens = nn.Embedding(self.vocab_size, self.hidden_size) # 语言模型 Decoder 层 self.layers = [ DecoderLayer( hidden_size=self.hidden_size, num_heads=self.num_attention_heads, intermediate_size=self.intermediate_size, num_kv_heads=self.num_kv_heads, head_dim=self.head_dim, ) for _ in range(self.num_hidden_layers) ] # 最终归一化 self.norm = RMSNorm(self.hidden_size, eps=config.get('rms_norm_eps', 1e-6)) # LM head self.lm_head = nn.Linear(self.hidden_size, self.vocab_size, bias=False) def encode_image(self, pixel_values: mx.array) -> mx.array: """编码图像 - 使用完整的视觉编码器""" # pixel_values: [B, H, W, C] = [B, 224, 224, 3] # MLX Conv2d 期望输入: [B, H, W, C] B, H, W, C = pixel_values.shape # Patch embedding x = self.patch_embedding(pixel_values) # [B, 16, 16, 1152] # 重塑为序列 x = x.reshape(B, -1, self.vision_hidden_size) # [B, 256, 1152] # 添加位置嵌入（只使用前 256 个位置） x = x + self.position_embedding[:256, :] # 通过视觉 Transformer 层 for layer in self.vision_layers: x = layer(x, None) # 归一化 x = self.vision_norm(x) # 视觉投影层 (mlp_AR) # 1. Pre-norm x = self.vision_pre_norm(x) # [B, 256, 1152] # 2. Spatial merge (2x2 -> 1) # 将 16x16 grid 合并成 8x8，hidden_size 从 1152 扩展到 4608 grid_size = 16 # sqrt(256) merge_size = 2 merged_grid_size = grid_size // merge_size # 8 # Reshape 到 grid x = x.reshape(B, grid_size, grid_size, self.vision_hidden_size) # [B, 16, 16, 1152] # Reshape 以进行合并 x = x.reshape( B, merged_grid_size, merge_size, merged_grid_size, merge_size, self.vision_hidden_size ) # [B, 8, 2, 8, 2, 1152] # 转置和合并 x = mx.transpose(x, (0, 1, 3, 2, 4, 5)) # [B, 8, 8, 2, 2, 1152] x = x.reshape(B, merged_grid_size * merged_grid_size, -1) # [B, 64, 4608] # 3. Linear layers x = self.vision_linear_1(x) # [B, 64, 4608] x = nn.gelu(x) x = self.vision_linear_2(x) # [B, 64, 1024] return x def forward(self, input_ids: mx.array, vision_embeds: Optional[mx.array] = None) -> mx.array: """前向传播""" text_embeds = self.embed_tokens(input_ids) if vision_embeds is not None: # 添加 vision start/end tokens # 配置中的 vision_start_token_id: 101305, vision_end_token_id: 101306 batch_size = vision_embeds.shape[0] # 创建 vision start token embedding vision_start_id = mx.array([[101305]]) # vision_start_token_id vision_start_embed = self.embed_tokens(vision_start_id) # [1, 1, 1024] # 创建 vision end token embedding vision_end_id = mx.array([[101306]]) # vision_end_token_id vision_end_embed = self.embed_tokens(vision_end_id) # [1, 1, 1024] # 构建序列: [vision_start] + [vision_embeds] + [vision_end] + [text] hidden_states = mx.concatenate([ vision_start_embed, vision_embeds, vision_end_embed, text_embeds ], axis=1) else: hidden_states = text_embeds for layer in self.layers: hidden_states = layer(hidden_states, None) hidden_states = self.norm(hidden_states) logits = self.lm_head(hidden_states) return logits model = ImprovedModel(self.config) print("✅ 模型创建完成") return model def _load_all_weights(self): """加载所有权重（包括视觉编码器）""" print("\n" + "="*60) print("🔄 加载所有权重...") print("="*60) weights_path = self.model_dir / "paddleocr_vl_mlx.npz" # 加载权重文件 print(f"\n📂 加载权重文件...") weights = mx.load(str(weights_path)) print(f"✅ 加载了 {len(weights)} 个权重张量") # 直接设置参数 print(f"\n🔧 加载权重...") loaded_count = 0 try: # 1. 视觉编码器权重 print(f"\n📸 加载视觉编码器权重...") # Patch embedding - 需要转置权重 if 'visual.vision_model.embeddings.patch_embedding.weight' in weights: # PyTorch: [out, in, h, w] -> MLX: [out, h, w, in] w = weights['visual.vision_model.embeddings.patch_embedding.weight'] w_transposed = mx.transpose(w, (0, 2, 3, 1)) # [1152, 14, 14, 3] self.model.patch_embedding.weight = w_transposed loaded_count += 1 print(f" ✅ Patch embedding weight: {w.shape} -> {w_transposed.shape}") if 'visual.vision_model.embeddings.patch_embedding.bias' in weights: self.model.patch_embedding.bias = weights['visual.vision_model.embeddings.patch_embedding.bias'] loaded_count += 1 # 位置嵌入 if 'visual.vision_model.embeddings.position_embedding.weight' in weights: self.model.position_embedding = weights['visual.vision_model.embeddings.position_embedding.weight'] loaded_count += 1 print(f" ✅ Position embedding: {weights['visual.vision_model.embeddings.position_embedding.weight'].shape}") # 视觉 Transformer 层（加载全部 27 层） num_vision_layers_to_load = min(27, len(self.model.vision_layers)) print(f" 加载全部 {num_vision_layers_to_load} 个视觉层...") for i in range(num_vision_layers_to_load): layer = self.model.vision_layers[i] prefix = f'visual.vision_model.encoder.layers.{i}' # 注意力层权重 for proj_name in ['q_proj', 'k_proj', 'v_proj']: w_key = f'{prefix}.self_attn.{proj_name}.weight' b_key = f'{prefix}.self_attn.{proj_name}.bias' if w_key in weights: proj = getattr(layer.self_attn, proj_name) proj.weight = weights[w_key] if b_key in weights: proj.bias = weights[b_key] loaded_count += 1 # o_proj (out_proj) w_key = f'{prefix}.self_attn.out_proj.weight' b_key = f'{prefix}.self_attn.out_proj.bias' if w_key in weights: layer.self_attn.o_proj.weight = weights[w_key] if b_key in weights: layer.self_attn.o_proj.bias = weights[b_key] loaded_count += 1 # MLP 层 if f'{prefix}.mlp.fc1.weight' in weights: layer.mlp.gate_proj.weight = weights[f'{prefix}.mlp.fc1.weight'] loaded_count += 1 if f'{prefix}.mlp.fc2.weight' in weights: layer.mlp.down_proj.weight = weights[f'{prefix}.mlp.fc2.weight'] loaded_count += 1 # LayerNorm for norm_name, model_norm in [('layer_norm1', 'input_layernorm'), ('layer_norm2', 'post_attention_layernorm')]: if f'{prefix}.{norm_name}.weight' in weights: getattr(layer, model_norm).weight = weights[f'{prefix}.{norm_name}.weight'] loaded_count += 1 print(f"✅ 视觉编码器权重加载完成 ({num_vision_layers_to_load} 层)") # 视觉投影层 (mlp_AR) print(f"\n🔗 加载视觉投影层 (mlp_AR)...") mlp_ar_loaded = 0 if 'mlp_AR.pre_norm.weight' in weights: self.model.vision_pre_norm.weight = weights['mlp_AR.pre_norm.weight'] mlp_ar_loaded += 1 if 'mlp_AR.pre_norm.bias' in weights: self.model.vision_pre_norm.bias = weights['mlp_AR.pre_norm.bias'] mlp_ar_loaded += 1 if 'mlp_AR.linear_1.weight' in weights: self.model.vision_linear_1.weight = weights['mlp_AR.linear_1.weight'] mlp_ar_loaded += 1 if 'mlp_AR.linear_1.bias' in weights: self.model.vision_linear_1.bias = weights['mlp_AR.linear_1.bias'] mlp_ar_loaded += 1 if 'mlp_AR.linear_2.weight' in weights: self.model.vision_linear_2.weight = weights['mlp_AR.linear_2.weight'] mlp_ar_loaded += 1 if 'mlp_AR.linear_2.bias' in weights: self.model.vision_linear_2.bias = weights['mlp_AR.linear_2.bias'] mlp_ar_loaded += 1 print(f"✅ 视觉投影层加载完成 ({mlp_ar_loaded}/6 个权重)") loaded_count += mlp_ar_loaded # 2. 语言模型权重（与之前相同） print(f"\n📝 加载语言模型权重...") # Token 嵌入 if 'model.embed_tokens.weight' in weights: self.model.embed_tokens.weight = weights['model.embed_tokens.weight'] loaded_count += 1 # 语言模型层 for i in range(18): layer = self.model.layers[i] prefix = f'model.layers.{i}' # 注意力层 for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']: if f'{prefix}.self_attn.{proj}.weight' in weights: getattr(layer.self_attn, proj).weight = weights[f'{prefix}.self_attn.{proj}.weight'] loaded_count += 1 # MLP 层 for proj in ['gate_proj', 'up_proj', 'down_proj']: if f'{prefix}.mlp.{proj}.weight' in weights: getattr(layer.mlp, proj).weight = weights[f'{prefix}.mlp.{proj}.weight'] loaded_count += 1 # 归一化层 for norm in ['input_layernorm', 'post_attention_layernorm']: if f'{prefix}.{norm}.weight' in weights: getattr(layer, norm).weight = weights[f'{prefix}.{norm}.weight'] loaded_count += 1 # 最终归一化 if 'model.norm.weight' in weights: self.model.norm.weight = weights['model.norm.weight'] loaded_count += 1 # LM head if 'lm_head.weight' in weights: self.model.lm_head.weight = weights['lm_head.weight'] loaded_count += 1 print(f"✅ 语言模型权重加载完成") print(f"\n✅ 总共成功加载 {loaded_count} 个权重") except Exception as e: print(f"\n❌ 权重加载失败: {e}") import traceback traceback.print_exc() print("\n" + "="*60) print(f"📊 权重加载完成") print("="*60) def preprocess_image(self, image_path: str) -> Tuple[mx.array, tuple]: """预处理图像""" image = Image.open(image_path).convert('RGB') original_size = image.size target_size = (224, 224) image = image.resize(target_size, Image.Resampling.BILINEAR) image_array = np.array(image).astype(np.float32) / 255.0 mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) image_array = (image_array - mean) / std image_array = np.expand_dims(image_array, 0) return mx.array(image_array), original_size def encode_prompt(self, prompt: str) -> mx.array: """编码提示文本""" if self.tokenizer: tokens = self.tokenizer.encode(prompt, add_special_tokens=True) return mx.array([tokens]) else: return mx.array([[1, 2, 3, 4, 5]]) def decode_tokens(self, token_ids: List[int]) -> str: """解码 token IDs 为文本""" if self.tokenizer: text = self.tokenizer.decode(token_ids, skip_special_tokens=True) return text else: return f"[Token IDs: {token_ids[:10]}...]" def generate( self, pixel_values: mx.array, prompt: str = "Extract all text from this image.", max_tokens: int = 100, temperature: float = 0.0, repetition_penalty: float = 1.2, ) -> str: """生成文本""" print(f"\n🔮 开始生成...") # 编码图像 start = time.time() vision_embeds = self.model.encode_image(pixel_values) print(f"✅ 图像编码: {vision_embeds.shape} ({time.time()-start:.2f}s)") # 编码提示 start = time.time() prompt_ids = self.encode_prompt(prompt) print(f"✅ 提示编码: {prompt_ids.shape} ({time.time()-start:.2f}s)") # 自回归生成 print(f"\n🔄 自回归生成 (max_tokens={max_tokens}, repetition_penalty={repetition_penalty})...") start = time.time() output_ids = [] current_ids = prompt_ids eos_token_id = self.tokenizer.eos_token_id if self.tokenizer else 2 for i in range(max_tokens): logits = self.model.forward(current_ids, vision_embeds) next_token_logits = logits[:, -1, :] # MLX 会自动创建新的视图 # 应用 repetition penalty if repetition_penalty != 1.0 and len(output_ids) > 0: # 创建一个可修改的副本 next_token_logits = mx.array(next_token_logits) for token_id in set(output_ids): # 使用 set 避免重复惩罚 next_token_logits[0, token_id] = next_token_logits[0, token_id] / repetition_penalty if temperature == 0: next_token = mx.argmax(next_token_logits, axis=-1) else: next_token_logits = next_token_logits / temperature probs = mx.softmax(next_token_logits, axis=-1) next_token = mx.random.categorical(probs) next_token_id = int(next_token[0]) if next_token_id == eos_token_id: print(f" 遇到 EOS token，停止生成") break output_ids.append(next_token_id) current_ids = mx.concatenate([current_ids, mx.array([[next_token_id]])], axis=1) if (i + 1) % 20 == 0: print(f" 生成了 {i + 1} tokens...") elapsed = time.time() - start print(f"✅ 生成完成: {len(output_ids)} tokens ({elapsed:.2f}s, {len(output_ids)/elapsed:.1f} tokens/s)") # 解码 result_text = self.decode_tokens(output_ids) return result_text def ocr( self, image_path: str, prompt: str = "Extract all text from this image.", max_tokens: int = 100, repetition_penalty: float = 1.2, ) -> dict: """端到端 OCR""" print("\n" + "="*60) print("🚀 执行改进版 OCR") print("="*60) total_start = time.time() # 预处理 pixel_values, original_size = self.preprocess_image(image_path) print(f"📸 图像: {original_size} -> {pixel_values.shape}") # 生成 result_text = self.generate(pixel_values, prompt, max_tokens, repetition_penalty=repetition_penalty) total_time = time.time() - total_start print(f"\n✅ OCR 完成 (总耗时: {total_time:.2f}s)") print("="*60) return { 'text': result_text, 'image_size': original_size, 'elapsed_time': total_time, 'status': 'success' } def main(): """主函数""" print("\n" + "="*60) print("🎯 PaddleOCR MLX 改进版测试") print("="*60) print(f"目标: 提升准确度") print(f"改进: 加载完整视觉编码器权重") print("="*60) model_dir = "/Users/gt/.gemini/antigravity/scratch/paddleocr-mlx-conversion" try: # 初始化 ocr = ImprovedPaddleOCRMLX(model_dir) # 创建测试图像 print("\n📋 创建测试图像...") img = Image.new('RGB', (400, 200), color='white') draw = ImageDraw.Draw(img) draw.text((50, 80), "Hello MLX!", fill='black') test_path = "/tmp/test_improved_mlx.png" img.save(test_path) print(f"✅ 测试图像: {test_path}") # 运行 OCR result = ocr.ocr(test_path, max_tokens=50) # 显示结果 print(f"\n📝 OCR 结果:") print(f"{'='*60}") print(result['text']) print(f"{'='*60}") print(f"图像大小: {result['image_size']}") print(f"耗时: {result['elapsed_time']:.2f}s") print(f"\n🎉 改进版测试完成！") except Exception as e: print(f"\n❌ 错误: {e}") import traceback traceback.print_exc() if __name__ == "__main__": main()