Add VLM picture description support for image handling
This commit is contained in:
parent
9bf3a83b5d
commit
3c4116397d
4 changed files with 194 additions and 1 deletions
|
|
@ -36,6 +36,13 @@ processing:
|
||||||
# Image settings
|
# Image settings
|
||||||
images_scale: 2.0 # Image scale factor
|
images_scale: 2.0 # Image scale factor
|
||||||
generate_picture_images: false # Include embedded images in output
|
generate_picture_images: false # Include embedded images in output
|
||||||
|
|
||||||
|
# VLM picture description (optional)
|
||||||
|
picture_description:
|
||||||
|
enabled: false # Enable VLM image descriptions
|
||||||
|
model:
|
||||||
|
provider: ollama
|
||||||
|
name: ministral-3
|
||||||
```
|
```
|
||||||
|
|
||||||
### Conversion Options
|
### Conversion Options
|
||||||
|
|
@ -81,6 +88,70 @@ conversion_options:
|
||||||
- **images_scale**: Scale factor for extracted images. Higher values = better quality but larger size. Typical range: 1.0-3.0.
|
- **images_scale**: Scale factor for extracted images. Higher values = better quality but larger size. Typical range: 1.0-3.0.
|
||||||
- **generate_picture_images**: When `true`, embedded images (figures, diagrams) are included as base64-encoded data in the document. When `false` (default), images are excluded to reduce chunk size and avoid context bloat.
|
- **generate_picture_images**: When `true`, embedded images (figures, diagrams) are included as base64-encoded data in the document. When `false` (default), images are excluded to reduce chunk size and avoid context bloat.
|
||||||
|
|
||||||
|
#### Picture Description (VLM)
|
||||||
|
|
||||||
|
Use a Vision Language Model (VLM) to automatically describe images in documents. Descriptions become searchable text, improving RAG retrieval for visual content.
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
conversion_options:
|
||||||
|
picture_description:
|
||||||
|
enabled: true # Enable VLM picture description
|
||||||
|
model:
|
||||||
|
provider: ollama # ollama, openai, or custom
|
||||||
|
name: ministral-3 # VLM model name
|
||||||
|
prompt: "Describe this image in detail. Be precise and concise."
|
||||||
|
timeout: 90 # Request timeout in seconds
|
||||||
|
max_tokens: 200 # Maximum tokens in response
|
||||||
|
```
|
||||||
|
|
||||||
|
**Configuration options:**
|
||||||
|
|
||||||
|
- **enabled**: When `true`, each embedded image is sent to a VLM for description. Requires `generate_picture_images` to be `true` (automatically enabled).
|
||||||
|
- **model**: Standard model configuration
|
||||||
|
- `provider`: `ollama` (default), `openai`, or use `base_url` for custom endpoints
|
||||||
|
- `name`: Model name (e.g., `ministral-3`, `granite3.2-vision`, `gpt-4-vision`)
|
||||||
|
- `base_url`: Optional custom API endpoint for vLLM, LM Studio, etc.
|
||||||
|
- **prompt**: Instruction for the VLM when describing images
|
||||||
|
- **timeout**: Request timeout in seconds
|
||||||
|
- **max_tokens**: Maximum tokens in the VLM response
|
||||||
|
|
||||||
|
**Using with Ollama:**
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
conversion_options:
|
||||||
|
picture_description:
|
||||||
|
enabled: true
|
||||||
|
model:
|
||||||
|
provider: ollama
|
||||||
|
name: ministral-3
|
||||||
|
```
|
||||||
|
|
||||||
|
Requires Ollama running with a vision-capable model:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
ollama pull ministral-3
|
||||||
|
ollama serve
|
||||||
|
```
|
||||||
|
|
||||||
|
**Using with vLLM or custom endpoints:**
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
conversion_options:
|
||||||
|
picture_description:
|
||||||
|
enabled: true
|
||||||
|
model:
|
||||||
|
provider: openai # Use OpenAI-compatible API format
|
||||||
|
name: granite-vision
|
||||||
|
base_url: http://my-vllm-server:8000
|
||||||
|
```
|
||||||
|
|
||||||
|
**How it works:**
|
||||||
|
|
||||||
|
1. During PDF conversion, docling extracts embedded images
|
||||||
|
2. Each image is sent to the configured VLM for description
|
||||||
|
3. Descriptions are added as annotations on the image
|
||||||
|
4. When exported to markdown, descriptions appear as searchable text
|
||||||
|
|
||||||
### Local vs Remote Processing
|
### Local vs Remote Processing
|
||||||
|
|
||||||
**Local processing** (default):
|
**Local processing** (default):
|
||||||
|
|
|
||||||
|
|
@ -96,6 +96,21 @@ class ResearchConfig(BaseModel):
|
||||||
max_concurrency: int = 1
|
max_concurrency: int = 1
|
||||||
|
|
||||||
|
|
||||||
|
class PictureDescriptionConfig(BaseModel):
|
||||||
|
"""Configuration for VLM-based picture description."""
|
||||||
|
|
||||||
|
enabled: bool = False
|
||||||
|
model: ModelConfig = Field(
|
||||||
|
default_factory=lambda: ModelConfig(
|
||||||
|
provider="ollama",
|
||||||
|
name="ministral-3",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
prompt: str = "Describe this image in detail. Be precise and concise."
|
||||||
|
timeout: int = 90
|
||||||
|
max_tokens: int = 200
|
||||||
|
|
||||||
|
|
||||||
class ConversionOptions(BaseModel):
|
class ConversionOptions(BaseModel):
|
||||||
"""Options for document conversion."""
|
"""Options for document conversion."""
|
||||||
|
|
||||||
|
|
@ -113,6 +128,11 @@ class ConversionOptions(BaseModel):
|
||||||
images_scale: float = 2.0
|
images_scale: float = 2.0
|
||||||
generate_picture_images: bool = False
|
generate_picture_images: bool = False
|
||||||
|
|
||||||
|
# VLM picture description
|
||||||
|
picture_description: PictureDescriptionConfig = Field(
|
||||||
|
default_factory=PictureDescriptionConfig
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class ProcessingConfig(BaseModel):
|
class ProcessingConfig(BaseModel):
|
||||||
chunk_size: int = 256
|
chunk_size: int = 256
|
||||||
|
|
|
||||||
|
|
@ -11,6 +11,8 @@ from haiku.rag.converters.text_utils import TextFileHandler
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from docling_core.types.doc.document import DoclingDocument
|
from docling_core.types.doc.document import DoclingDocument
|
||||||
|
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
|
||||||
class DoclingLocalConverter(DocumentConverter):
|
class DoclingLocalConverter(DocumentConverter):
|
||||||
"""Converter that uses local docling for document conversion.
|
"""Converter that uses local docling for document conversion.
|
||||||
|
|
@ -54,6 +56,21 @@ class DoclingLocalConverter(DocumentConverter):
|
||||||
"""Return list of file extensions supported by this converter."""
|
"""Return list of file extensions supported by this converter."""
|
||||||
return self.docling_extensions + TextFileHandler.text_extensions
|
return self.docling_extensions + TextFileHandler.text_extensions
|
||||||
|
|
||||||
|
def _get_vlm_api_url(self, model: "ModelConfig") -> str:
|
||||||
|
"""Construct VLM API URL from model config."""
|
||||||
|
if model.base_url:
|
||||||
|
base = model.base_url.rstrip("/")
|
||||||
|
return f"{base}/v1/chat/completions"
|
||||||
|
|
||||||
|
if model.provider == "ollama":
|
||||||
|
base = self.config.providers.ollama.base_url.rstrip("/")
|
||||||
|
return f"{base}/v1/chat/completions"
|
||||||
|
|
||||||
|
if model.provider == "openai":
|
||||||
|
return "https://api.openai.com/v1/chat/completions"
|
||||||
|
|
||||||
|
raise ValueError(f"Unsupported VLM provider: {model.provider}")
|
||||||
|
|
||||||
def _sync_convert_docling_file(self, path: Path) -> "DoclingDocument":
|
def _sync_convert_docling_file(self, path: Path) -> "DoclingDocument":
|
||||||
"""Synchronous conversion of docling-supported files."""
|
"""Synchronous conversion of docling-supported files."""
|
||||||
from docling.backend.docling_parse_backend import DoclingParseDocumentBackend
|
from docling.backend.docling_parse_backend import DoclingParseDocumentBackend
|
||||||
|
|
@ -61,6 +78,7 @@ class DoclingLocalConverter(DocumentConverter):
|
||||||
from docling.datamodel.pipeline_options import (
|
from docling.datamodel.pipeline_options import (
|
||||||
OcrAutoOptions,
|
OcrAutoOptions,
|
||||||
PdfPipelineOptions,
|
PdfPipelineOptions,
|
||||||
|
PictureDescriptionApiOptions,
|
||||||
TableFormerMode,
|
TableFormerMode,
|
||||||
TableStructureOptions,
|
TableStructureOptions,
|
||||||
)
|
)
|
||||||
|
|
@ -73,13 +91,14 @@ class DoclingLocalConverter(DocumentConverter):
|
||||||
)
|
)
|
||||||
|
|
||||||
opts = self.config.processing.conversion_options
|
opts = self.config.processing.conversion_options
|
||||||
|
pic_desc = opts.picture_description
|
||||||
|
|
||||||
pipeline_options = PdfPipelineOptions(
|
pipeline_options = PdfPipelineOptions(
|
||||||
do_ocr=opts.do_ocr,
|
do_ocr=opts.do_ocr,
|
||||||
do_table_structure=opts.do_table_structure,
|
do_table_structure=opts.do_table_structure,
|
||||||
images_scale=opts.images_scale,
|
images_scale=opts.images_scale,
|
||||||
generate_page_images=True,
|
generate_page_images=True,
|
||||||
generate_picture_images=opts.generate_picture_images,
|
generate_picture_images=opts.generate_picture_images or pic_desc.enabled,
|
||||||
table_structure_options=TableStructureOptions(
|
table_structure_options=TableStructureOptions(
|
||||||
do_cell_matching=opts.table_cell_matching,
|
do_cell_matching=opts.table_cell_matching,
|
||||||
mode=(
|
mode=(
|
||||||
|
|
@ -92,6 +111,21 @@ class DoclingLocalConverter(DocumentConverter):
|
||||||
force_full_page_ocr=opts.force_ocr,
|
force_full_page_ocr=opts.force_ocr,
|
||||||
lang=opts.ocr_lang if opts.ocr_lang else [],
|
lang=opts.ocr_lang if opts.ocr_lang else [],
|
||||||
),
|
),
|
||||||
|
do_picture_description=pic_desc.enabled,
|
||||||
|
)
|
||||||
|
|
||||||
|
if pic_desc.enabled:
|
||||||
|
from pydantic import AnyUrl
|
||||||
|
|
||||||
|
pipeline_options.enable_remote_services = True
|
||||||
|
pipeline_options.picture_description_options = PictureDescriptionApiOptions(
|
||||||
|
url=AnyUrl(self._get_vlm_api_url(pic_desc.model)),
|
||||||
|
params=dict(
|
||||||
|
model=pic_desc.model.name,
|
||||||
|
max_completion_tokens=pic_desc.max_tokens,
|
||||||
|
),
|
||||||
|
prompt=pic_desc.prompt,
|
||||||
|
timeout=pic_desc.timeout,
|
||||||
)
|
)
|
||||||
|
|
||||||
format_options = cast(
|
format_options = cast(
|
||||||
|
|
|
||||||
|
|
@ -308,6 +308,74 @@ class TestDoclingLocalConverter:
|
||||||
"Pictures should have image data when generate_picture_images=True"
|
"Pictures should have image data when generate_picture_images=True"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def test_get_vlm_api_url_with_ollama(self, config):
|
||||||
|
"""Test VLM API URL construction for Ollama provider."""
|
||||||
|
converter = DoclingLocalConverter(config)
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
model = ModelConfig(provider="ollama", name="ministral-3")
|
||||||
|
url = converter._get_vlm_api_url(model)
|
||||||
|
assert url == "http://localhost:11434/v1/chat/completions"
|
||||||
|
|
||||||
|
def test_get_vlm_api_url_with_custom_base_url(self, config):
|
||||||
|
"""Test VLM API URL construction with custom base_url."""
|
||||||
|
converter = DoclingLocalConverter(config)
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
model = ModelConfig(
|
||||||
|
provider="openai", name="gpt-4-vision", base_url="http://my-vllm:8000"
|
||||||
|
)
|
||||||
|
url = converter._get_vlm_api_url(model)
|
||||||
|
assert url == "http://my-vllm:8000/v1/chat/completions"
|
||||||
|
|
||||||
|
def test_get_vlm_api_url_with_openai(self, config):
|
||||||
|
"""Test VLM API URL construction for OpenAI provider."""
|
||||||
|
converter = DoclingLocalConverter(config)
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
model = ModelConfig(provider="openai", name="gpt-4-vision")
|
||||||
|
url = converter._get_vlm_api_url(model)
|
||||||
|
assert url == "https://api.openai.com/v1/chat/completions"
|
||||||
|
|
||||||
|
def test_get_vlm_api_url_unsupported_provider(self, config):
|
||||||
|
"""Test VLM API URL construction raises error for unsupported provider."""
|
||||||
|
converter = DoclingLocalConverter(config)
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
model = ModelConfig(provider="unsupported", name="test")
|
||||||
|
with pytest.raises(ValueError, match="Unsupported VLM provider"):
|
||||||
|
converter._get_vlm_api_url(model)
|
||||||
|
|
||||||
|
def test_picture_description_config_defaults(self, config):
|
||||||
|
"""Test that picture description config has correct defaults."""
|
||||||
|
assert config.processing.conversion_options.picture_description.enabled is False
|
||||||
|
assert (
|
||||||
|
config.processing.conversion_options.picture_description.model.provider
|
||||||
|
== "ollama"
|
||||||
|
)
|
||||||
|
assert (
|
||||||
|
config.processing.conversion_options.picture_description.model.name
|
||||||
|
== "ministral-3"
|
||||||
|
)
|
||||||
|
assert config.processing.conversion_options.picture_description.timeout == 90
|
||||||
|
assert (
|
||||||
|
config.processing.conversion_options.picture_description.max_tokens == 200
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_picture_description_config_applied(self, config):
|
||||||
|
"""Test that picture description config is applied to converter."""
|
||||||
|
config.processing.conversion_options.picture_description.enabled = True
|
||||||
|
config.processing.conversion_options.picture_description.prompt = (
|
||||||
|
"Custom prompt for testing."
|
||||||
|
)
|
||||||
|
config.processing.conversion_options.picture_description.timeout = 120
|
||||||
|
converter = DoclingLocalConverter(config)
|
||||||
|
|
||||||
|
pic_desc = converter.config.processing.conversion_options.picture_description
|
||||||
|
assert pic_desc.enabled is True
|
||||||
|
assert pic_desc.prompt == "Custom prompt for testing."
|
||||||
|
assert pic_desc.timeout == 120
|
||||||
|
|
||||||
|
|
||||||
class TestDoclingServeConverter:
|
class TestDoclingServeConverter:
|
||||||
"""Tests for DoclingServeConverter (mocked)."""
|
"""Tests for DoclingServeConverter (mocked)."""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue