[OCR] Phase 2-2: 多模型支持 (InternVL & LLaVA)
## 📋 任务概述
扩展 VLM-OCR 模块的模型支持，添加 InternVL 和 LLaVA 模型，实现多模型统一接口。
**父任务**: #2348 Phase 2: 功能完善
---
## 🎯 实现目标
### 1. 模型支持
- [ ] **InternVL-2.5**: 通用多模态理解模型
- [ ] **LLaVA-1.6**: 视觉指令遵循模型
- [ ] 统一的模型接口和配置
### 2. 模型选择
- [ ] API 端点支持模型参数
- [ ] 配置文件模型切换
- [ ] 模型性能对比工具
### 3. 模型管理
- [ ] 多模型同时加载（可选）
- [ ] 模型热切换
- [ ] 模型下载和缓存
---
## 📦 实现细节
### 2.1 InternVL 模型封装
```python
# mindnlp-ocr/models/internvl.py
from transformers import AutoModel, AutoTokenizer
from .base import VLMModelBase
class InternVLModel(VLMModelBase):
"""InternVL 模型封装"""
MODEL_NAME = "OpenGVLab/InternVL2_5-8B"
def __init__(self, model_name: str = None, device: str = "cuda"):
super().__init__(model_name or self.MODEL_NAME, device)
def load_model(self):
"""加载 InternVL 模型"""
self.model = AutoModel.from_pretrained(
self.model_name,
trust_remote_code=True,
torch_dtype=torch.float16
).to(self.device).eval()
self.tokenizer = AutoTokenizer.from_pretrained(
self.model_name,
trust_remote_code=True
)
def preprocess(self, image: Image, prompt: str) -> Dict:
"""InternVL 预处理"""
pixel_values = self.image_processor(
images=image,
return_tensors="pt"
)["pixel_values"].to(self.device)
input_ids = self.tokenizer.encode(
prompt,
return_tensors="pt"
).to(self.device)
return {
"pixel_values": pixel_values,
"input_ids": input_ids
}
def generate(self, inputs: Dict, **kwargs) -> str:
"""生成文本"""
with torch.no_grad():
outputs = self.model.generate(
pixel_values=inputs["pixel_values"],
input_ids=inputs["input_ids"],
max_new_tokens=kwargs.get("max_new_tokens", 512),
do_sample=kwargs.get("do_sample", False),
temperature=kwargs.get("temperature", 1.0)
)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
```
### 2.2 LLaVA 模型封装
```python
# mindnlp-ocr/models/llava.py
from transformers import LlavaForConditionalGeneration, AutoProcessor
from .base import VLMModelBase
class LLaVAModel(VLMModelBase):
"""LLaVA 模型封装"""
MODEL_NAME = "llava-hf/llava-1.5-7b-hf"
def __init__(self, model_name: str = None, device: str = "cuda"):
super().__init__(model_name or self.MODEL_NAME, device)
def load_model(self):
"""加载 LLaVA 模型"""
self.model = LlavaForConditionalGeneration.from_pretrained(
self.model_name,
torch_dtype=torch.float16,
low_cpu_mem_usage=True
).to(self.device)
self.processor = AutoProcessor.from_pretrained(self.model_name)
def preprocess(self, image: Image, prompt: str) -> Dict:
"""LLaVA 预处理"""
# LLaVA 需要特殊的 prompt 格式
formatted_prompt = f"USER: <image>\\n{prompt}\\nASSISTANT:"
inputs = self.processor(
text=formatted_prompt,
images=image,
return_tensors="pt"
).to(self.device)
return inputs
def generate(self, inputs: Dict, **kwargs) -> str:
"""生成文本"""
with torch.no_grad():
outputs = self.model.generate(
**inputs,
max_new_tokens=kwargs.get("max_new_tokens", 512),
do_sample=kwargs.get("do_sample", False)
)
return self.processor.decode(outputs[0], skip_special_tokens=True)
```
### 2.3 模型工厂和加载器
```python
# mindnlp-ocr/models/loader.py
from typing import Union
from .qwen2vl import Qwen2VLModel
from .internvl import InternVLModel
from .llava import LLaVAModel
from .base import VLMModelBase
class ModelFactory:
"""模型工厂"""
MODELS = {
"qwen2-vl": Qwen2VLModel,
"internvl": InternVLModel,
"llava": LLaVAModel
}
@classmethod
def create_model(cls, model_type: str, **kwargs) -> VLMModelBase:
"""创建模型实例"""
if model_type not in cls.MODELS:
raise ValueError(f"不支持的模型类型: {model_type}")
model_class = cls.MODELS[model_type]
return model_class(**kwargs)
@classmethod
def list_models(cls) -> list:
"""列出所有支持的模型"""
return list(cls.MODELS.keys())
class MultiModelLoader:
"""多模型加载器（支持同时加载多个模型）"""
def __init__(self):
self.models = {}
def load(self, model_name: str, model_type: str, **kwargs):
"""加载模型"""
if model_name in self.models:
return self.models[model_name]
model = ModelFactory.create_model(model_type, **kwargs)
model.load_model()
self.models[model_name] = model
return model
def get(self, model_name: str) -> VLMModelBase:
"""获取已加载的模型"""
return self.models.get(model_name)
def unload(self, model_name: str):
"""卸载模型"""
if model_name in self.models:
del self.models[model_name]
torch.cuda.empty_cache()
```
### 2.4 配置管理更新
```python
# mindnlp-ocr/config/settings.py
class Settings(BaseSettings):
# ... 现有配置 ...
# 模型配置
default_model: str = Field(
default="qwen2-vl",
description="默认使用的模型"
)
available_models: Dict[str, str] = Field(
default={
"qwen2-vl": "Qwen/Qwen2-VL-7B-Instruct",
"internvl": "OpenGVLab/InternVL2_5-8B",
"llava": "llava-hf/llava-1.5-7b-hf"
},
description="可用的模型列表"
)
enable_multi_model: bool = Field(
default=False,
description="是否启用多模型同时加载"
)
```
### 2.5 API 端点更新
```python
# mindnlp-ocr/api/routes/ocr.py
@router.post("/predict")
async def predict(
file: UploadFile = File(...),
prompt: str = Form(None),
model: str = Form("qwen2-vl"), # 新增模型参数
output_format: str = Form("json")
):
"""OCR 预测（支持模型选择）"""
# 验证模型
if model not in settings.available_models:
raise HTTPException(
status_code=400,
detail=f"不支持的模型: {model}。可用模型: {list(settings.available_models.keys())}"
)
# 使用指定模型
engine = get_engine(model)
result = engine.predict(image, prompt, output_format)
return result
@router.get("/models")
async def list_models():
"""列出所有可用的模型"""
return {
"models": [
{
"name": name,
"model_id": model_id,
"description": MODEL_DESCRIPTIONS.get(name, "")
}
for name, model_id in settings.available_models.items()
],
"default": settings.default_model
}
```
---
## 🧪 测试要求
### 功能测试
```bash
pytest tests/test_multi_models.py -v
```
### 测试用例
- [ ] `test_internvl_model`: InternVL 模型推理
- [ ] `test_llava_model`: LLaVA 模型推理
- [ ] `test_model_factory`: 模型工厂创建
- [ ] `test_model_switching`: 模型切换
- [ ] `test_multi_model_loading`: 多模型同时加载
- [ ] `test_model_comparison`: 模型性能对比
### 性能对比测试
```python
# tests/test_model_comparison.py
def test_model_comparison():
"""对比不同模型的性能"""
models = ["qwen2-vl", "internvl", "llava"]
results = {}
for model in models:
engine = VLMOCREngine(model_type=model)
# 测试准确率
accuracy = evaluate_accuracy(engine, test_dataset)
# 测试速度
speed = evaluate_speed(engine, test_images)
results[model] = {
"accuracy": accuracy,
"speed": speed,
"memory": get_gpu_memory()
}
print(f"模型对比: {results}")
```
---
## 📊 模型特性对比
| 模型 | 参数量 | 输入分辨率 | 推理速度 | 优势 | 适用场景 |
|------|--------|-----------|---------|------|---------|
| **Qwen2-VL** | 7B | 动态 | 快 | 多语言、高精度 | 通用OCR |
| **InternVL** | 8B | 448×448 | 中等 | 多模态理解强 | 文档分析 |
| **LLaVA** | 7B | 336×336 | 最快 | 指令遵循好 | 简单识别 |
---
## 📚 参考资料
- [InternVL GitHub](https://github.com/OpenGVLab/InternVL)
- [LLaVA HuggingFace](https://huggingface.co/llava-hf)
- [Model Comparison Guide](https://github.com/BradyFU/Awesome-Multimodal-Large-Language-Models)
---
## ✅ 验收标准
- [ ] InternVL 模型成功集成并通过测试
- [ ] LLaVA 模型成功集成并通过测试
- [ ] 模型工厂正常工作
- [ ] API 支持模型选择参数
- [ ] 模型性能对比文档完善
- [ ] 通过所有单元测试和集成测试
- [ ] 多模型使用文档完善
---
**优先级**: P1
**预计工作量**: 7-10天
**依赖**: #2351 (模型层基础完成)
0 条评论