mirror of
https://github.com/DS4SD/docling.git
synced 2025-07-26 20:14:47 +00:00
feat: adding new vlm-models support
Signed-off-by: Peter Staar <taa@zurich.ibm.com>
This commit is contained in:
parent
f1658edbad
commit
18e1ec4df2
@ -267,7 +267,8 @@ class InferenceFramework(str, Enum):
|
|||||||
MLX = "mlx"
|
MLX = "mlx"
|
||||||
TRANSFORMERS = "transformers"
|
TRANSFORMERS = "transformers"
|
||||||
OPENAI = "openai"
|
OPENAI = "openai"
|
||||||
|
TRANSFORMERS_AutoModelForVision2Seq = "transformers-AutoModelForVision2Seq"
|
||||||
|
TRANSFORMERS_AutoModelForCausalLM = "transformers-AutoModelForCausalLM"
|
||||||
|
|
||||||
class HuggingFaceVlmOptions(BaseVlmOptions):
|
class HuggingFaceVlmOptions(BaseVlmOptions):
|
||||||
kind: Literal["hf_model_options"] = "hf_model_options"
|
kind: Literal["hf_model_options"] = "hf_model_options"
|
||||||
@ -310,7 +311,7 @@ smoldocling_vlm_conversion_options = HuggingFaceVlmOptions(
|
|||||||
repo_id="ds4sd/SmolDocling-256M-preview",
|
repo_id="ds4sd/SmolDocling-256M-preview",
|
||||||
prompt="Convert this page to docling.",
|
prompt="Convert this page to docling.",
|
||||||
response_format=ResponseFormat.DOCTAGS,
|
response_format=ResponseFormat.DOCTAGS,
|
||||||
inference_framework=InferenceFramework.TRANSFORMERS,
|
inference_framework=InferenceFramework.TRANSFORMERS_AutoModelForVision2Seq,
|
||||||
)
|
)
|
||||||
|
|
||||||
granite_vision_vlm_conversion_options = HuggingFaceVlmOptions(
|
granite_vision_vlm_conversion_options = HuggingFaceVlmOptions(
|
||||||
@ -318,7 +319,7 @@ granite_vision_vlm_conversion_options = HuggingFaceVlmOptions(
|
|||||||
# prompt="OCR the full page to markdown.",
|
# prompt="OCR the full page to markdown.",
|
||||||
prompt="OCR this image.",
|
prompt="OCR this image.",
|
||||||
response_format=ResponseFormat.MARKDOWN,
|
response_format=ResponseFormat.MARKDOWN,
|
||||||
inference_framework=InferenceFramework.TRANSFORMERS,
|
inference_framework=InferenceFramework.TRANSFORMERS_AutoModelForVision2Seq,
|
||||||
)
|
)
|
||||||
|
|
||||||
granite_vision_vlm_ollama_conversion_options = ApiVlmOptions(
|
granite_vision_vlm_ollama_conversion_options = ApiVlmOptions(
|
||||||
|
@ -57,11 +57,14 @@ class HuggingFaceVlmModel(BasePageModel):
|
|||||||
)
|
)
|
||||||
self.param_quantized = vlm_options.quantized # False
|
self.param_quantized = vlm_options.quantized # False
|
||||||
|
|
||||||
self.processor = AutoProcessor.from_pretrained(artifacts_path)
|
self.processor = AutoProcessor.from_pretrained(
|
||||||
|
artifacts_path,
|
||||||
|
# trust_remote_code=True,
|
||||||
|
)
|
||||||
if not self.param_quantized:
|
if not self.param_quantized:
|
||||||
self.vlm_model = AutoModelForVision2Seq.from_pretrained(
|
self.vlm_model = AutoModelForVision2Seq.from_pretrained(
|
||||||
artifacts_path,
|
artifacts_path,
|
||||||
device_map=device,
|
device_map=self.device,
|
||||||
torch_dtype=torch.bfloat16,
|
torch_dtype=torch.bfloat16,
|
||||||
_attn_implementation=(
|
_attn_implementation=(
|
||||||
"flash_attention_2"
|
"flash_attention_2"
|
||||||
@ -69,12 +72,13 @@ class HuggingFaceVlmModel(BasePageModel):
|
|||||||
and accelerator_options.cuda_use_flash_attention2
|
and accelerator_options.cuda_use_flash_attention2
|
||||||
else "eager"
|
else "eager"
|
||||||
),
|
),
|
||||||
|
# trust_remote_code=True,
|
||||||
) # .to(self.device)
|
) # .to(self.device)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
self.vlm_model = AutoModelForVision2Seq.from_pretrained(
|
self.vlm_model = AutoModelForVision2Seq.from_pretrained(
|
||||||
artifacts_path,
|
artifacts_path,
|
||||||
device_map=device,
|
device_map=self.device,
|
||||||
torch_dtype="auto",
|
torch_dtype="auto",
|
||||||
quantization_config=self.param_quantization_config,
|
quantization_config=self.param_quantization_config,
|
||||||
_attn_implementation=(
|
_attn_implementation=(
|
||||||
@ -83,6 +87,7 @@ class HuggingFaceVlmModel(BasePageModel):
|
|||||||
and accelerator_options.cuda_use_flash_attention2
|
and accelerator_options.cuda_use_flash_attention2
|
||||||
else "eager"
|
else "eager"
|
||||||
),
|
),
|
||||||
|
# trust_remote_code=True,
|
||||||
) # .to(self.device)
|
) # .to(self.device)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
|
0
docling/models/hf_vlm_models/__init__.py
Normal file
0
docling/models/hf_vlm_models/__init__.py
Normal file
@ -0,0 +1,233 @@
|
|||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from collections.abc import Iterable
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import Page, VlmPrediction
|
||||||
|
from docling.datamodel.document import ConversionResult
|
||||||
|
from docling.datamodel.pipeline_options import (
|
||||||
|
AcceleratorOptions,
|
||||||
|
HuggingFaceVlmOptions,
|
||||||
|
)
|
||||||
|
from docling.models.base_model import BasePageModel
|
||||||
|
from docling.utils.accelerator_utils import decide_device
|
||||||
|
from docling.utils.profiling import TimeRecorder
|
||||||
|
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class HuggingFaceVlmModel_AutoModelForCausalLM(BasePageModel):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
enabled: bool,
|
||||||
|
artifacts_path: Optional[Path],
|
||||||
|
accelerator_options: AcceleratorOptions,
|
||||||
|
vlm_options: HuggingFaceVlmOptions,
|
||||||
|
):
|
||||||
|
self.enabled = enabled
|
||||||
|
|
||||||
|
self.trust_remote_code = True
|
||||||
|
|
||||||
|
self.vlm_options = vlm_options
|
||||||
|
|
||||||
|
if self.enabled:
|
||||||
|
import torch
|
||||||
|
from transformers import ( # type: ignore
|
||||||
|
AutoModelForCausalLM,
|
||||||
|
AutoProcessor,
|
||||||
|
GenerationConfig,
|
||||||
|
BitsAndBytesConfig,
|
||||||
|
)
|
||||||
|
|
||||||
|
device = decide_device(accelerator_options.device)
|
||||||
|
self.device = 'cpu' #device
|
||||||
|
|
||||||
|
_log.debug(f"Available device for HuggingFace VLM: {device}")
|
||||||
|
print(f"Available device for HuggingFace VLM: {device}")
|
||||||
|
|
||||||
|
repo_cache_folder = vlm_options.repo_id.replace("/", "--")
|
||||||
|
|
||||||
|
# PARAMETERS:
|
||||||
|
if artifacts_path is None:
|
||||||
|
artifacts_path = self.download_models(self.vlm_options.repo_id)
|
||||||
|
elif (artifacts_path / repo_cache_folder).exists():
|
||||||
|
artifacts_path = artifacts_path / repo_cache_folder
|
||||||
|
|
||||||
|
self.param_question = vlm_options.prompt # "Perform Layout Analysis."
|
||||||
|
self.param_quantization_config = BitsAndBytesConfig(
|
||||||
|
load_in_8bit=vlm_options.load_in_8bit, # True,
|
||||||
|
llm_int8_threshold=vlm_options.llm_int8_threshold, # 6.0
|
||||||
|
)
|
||||||
|
self.param_quantized = vlm_options.quantized # False
|
||||||
|
|
||||||
|
self.processor = AutoProcessor.from_pretrained(
|
||||||
|
artifacts_path,
|
||||||
|
trust_remote_code=self.trust_remote_code,
|
||||||
|
)
|
||||||
|
if not self.param_quantized:
|
||||||
|
self.vlm_model = AutoModelForCausalLM.from_pretrained(
|
||||||
|
artifacts_path,
|
||||||
|
device_map=self.device,
|
||||||
|
torch_dtype=torch.bfloat16,
|
||||||
|
_attn_implementation=(
|
||||||
|
"flash_attention_2"
|
||||||
|
if self.device.startswith("cuda")
|
||||||
|
and accelerator_options.cuda_use_flash_attention2
|
||||||
|
else "eager"
|
||||||
|
),
|
||||||
|
trust_remote_code=self.trust_remote_code,
|
||||||
|
).to(self.device)
|
||||||
|
|
||||||
|
else:
|
||||||
|
self.vlm_model = AutoModelForCausalLM.from_pretrained(
|
||||||
|
artifacts_path,
|
||||||
|
device_map=self.device,
|
||||||
|
torch_dtype="auto",
|
||||||
|
quantization_config=self.param_quantization_config,
|
||||||
|
_attn_implementation=(
|
||||||
|
"flash_attention_2"
|
||||||
|
if self.device.startswith("cuda")
|
||||||
|
and accelerator_options.cuda_use_flash_attention2
|
||||||
|
else "eager"
|
||||||
|
),
|
||||||
|
trust_remote_code=self.trust_remote_code,
|
||||||
|
).to(self.device)
|
||||||
|
|
||||||
|
model_path = artifacts_path
|
||||||
|
print(f"model: {model_path}")
|
||||||
|
|
||||||
|
# Load generation config
|
||||||
|
self.generation_config = GenerationConfig.from_pretrained(model_path)
|
||||||
|
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def download_models(
|
||||||
|
repo_id: str,
|
||||||
|
local_dir: Optional[Path] = None,
|
||||||
|
force: bool = False,
|
||||||
|
progress: bool = False,
|
||||||
|
) -> Path:
|
||||||
|
from huggingface_hub import snapshot_download
|
||||||
|
from huggingface_hub.utils import disable_progress_bars
|
||||||
|
|
||||||
|
if not progress:
|
||||||
|
disable_progress_bars()
|
||||||
|
download_path = snapshot_download(
|
||||||
|
repo_id=repo_id,
|
||||||
|
force_download=force,
|
||||||
|
local_dir=local_dir,
|
||||||
|
# revision="v0.0.1",
|
||||||
|
)
|
||||||
|
|
||||||
|
return Path(download_path)
|
||||||
|
|
||||||
|
def __call__(
|
||||||
|
self, conv_res: ConversionResult, page_batch: Iterable[Page]
|
||||||
|
) -> Iterable[Page]:
|
||||||
|
for page in page_batch:
|
||||||
|
assert page._backend is not None
|
||||||
|
if not page._backend.is_valid():
|
||||||
|
yield page
|
||||||
|
else:
|
||||||
|
with TimeRecorder(conv_res, "vlm"):
|
||||||
|
assert page.size is not None
|
||||||
|
|
||||||
|
# hi_res_image = page.get_image(scale=2.0) # 144dpi
|
||||||
|
hi_res_image = page.get_image(scale=1.0) # 72dpi
|
||||||
|
|
||||||
|
if hi_res_image is not None:
|
||||||
|
im_width, im_height = hi_res_image.size
|
||||||
|
|
||||||
|
# populate page_tags with predicted doc tags
|
||||||
|
page_tags = ""
|
||||||
|
|
||||||
|
if hi_res_image:
|
||||||
|
if hi_res_image.mode != "RGB":
|
||||||
|
hi_res_image = hi_res_image.convert("RGB")
|
||||||
|
|
||||||
|
"""
|
||||||
|
messages = [
|
||||||
|
{
|
||||||
|
"role": "user",
|
||||||
|
"content": [
|
||||||
|
{
|
||||||
|
"type": "text",
|
||||||
|
"text": "This is a page from a document.",
|
||||||
|
},
|
||||||
|
{"type": "image"},
|
||||||
|
{"type": "text", "text": self.param_question},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
]
|
||||||
|
prompt = self.processor.apply_chat_template(
|
||||||
|
messages, add_generation_prompt=False
|
||||||
|
)
|
||||||
|
inputs = self.processor(
|
||||||
|
text=prompt, images=[hi_res_image], return_tensors="pt"
|
||||||
|
)
|
||||||
|
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
||||||
|
|
||||||
|
start_time = time.time()
|
||||||
|
# Call model to generate:
|
||||||
|
generated_ids = self.vlm_model.generate(
|
||||||
|
**inputs, max_new_tokens=4096, use_cache=True
|
||||||
|
)
|
||||||
|
|
||||||
|
generation_time = time.time() - start_time
|
||||||
|
generated_texts = self.processor.batch_decode(
|
||||||
|
generated_ids[:, inputs["input_ids"].shape[1] :],
|
||||||
|
skip_special_tokens=False,
|
||||||
|
)[0]
|
||||||
|
|
||||||
|
num_tokens = len(generated_ids[0])
|
||||||
|
page_tags = generated_texts
|
||||||
|
"""
|
||||||
|
|
||||||
|
hi_res_image.show()
|
||||||
|
|
||||||
|
# Define prompt structure
|
||||||
|
user_prompt = '<|user|>'
|
||||||
|
assistant_prompt = '<|assistant|>'
|
||||||
|
prompt_suffix = '<|end|>'
|
||||||
|
|
||||||
|
# Part 1: Image Processing
|
||||||
|
print("\n--- IMAGE PROCESSING ---")
|
||||||
|
# image_url = 'https://www.ilankelman.org/stopsigns/australia.jpg'
|
||||||
|
prompt = f'{user_prompt}<|image_1|>OCR this image into MarkDown?{prompt_suffix}{assistant_prompt}'
|
||||||
|
print(f'>>> Prompt\n{prompt}')
|
||||||
|
|
||||||
|
inputs = self.processor(text=prompt, images=hi_res_image, return_tensors='pt').to(self.device) #.to('cuda:0')
|
||||||
|
print("inputs: ", inputs.keys())
|
||||||
|
|
||||||
|
# Generate response
|
||||||
|
generate_ids = self.vlm_model.generate(
|
||||||
|
**inputs,
|
||||||
|
max_new_tokens=128,
|
||||||
|
generation_config=self.generation_config,
|
||||||
|
)
|
||||||
|
generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
|
||||||
|
|
||||||
|
num_tokens = len(generated_ids[0])
|
||||||
|
response = self.processor.batch_decode(
|
||||||
|
generate_ids,
|
||||||
|
skip_special_tokens=True,
|
||||||
|
clean_up_tokenization_spaces=False
|
||||||
|
)[0]
|
||||||
|
print(f'>>> Response\n{response}')
|
||||||
|
|
||||||
|
_log.debug(
|
||||||
|
f"Generated {num_tokens} tokens in time {generation_time:.2f} seconds."
|
||||||
|
)
|
||||||
|
|
||||||
|
# inference_time = time.time() - start_time
|
||||||
|
# tokens_per_second = num_tokens / generation_time
|
||||||
|
# print("")
|
||||||
|
# print(f"Page Inference Time: {inference_time:.2f} seconds")
|
||||||
|
# print(f"Total tokens on page: {num_tokens:.2f}")
|
||||||
|
# print(f"Tokens/sec: {tokens_per_second:.2f}")
|
||||||
|
# print("")
|
||||||
|
page.predictions.vlm_response = VlmPrediction(text=response)
|
||||||
|
|
||||||
|
yield page
|
@ -0,0 +1,187 @@
|
|||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from collections.abc import Iterable
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import Page, VlmPrediction
|
||||||
|
from docling.datamodel.document import ConversionResult
|
||||||
|
from docling.datamodel.pipeline_options import (
|
||||||
|
AcceleratorOptions,
|
||||||
|
HuggingFaceVlmOptions,
|
||||||
|
)
|
||||||
|
from docling.models.base_model import BasePageModel
|
||||||
|
from docling.utils.accelerator_utils import decide_device
|
||||||
|
from docling.utils.profiling import TimeRecorder
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class HuggingFaceVlmModel_AutoModelForVision2Seq(BasePageModel):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
enabled: bool,
|
||||||
|
artifacts_path: Optional[Path],
|
||||||
|
accelerator_options: AcceleratorOptions,
|
||||||
|
vlm_options: HuggingFaceVlmOptions,
|
||||||
|
):
|
||||||
|
self.enabled = enabled
|
||||||
|
|
||||||
|
self.vlm_options = vlm_options
|
||||||
|
|
||||||
|
if self.enabled:
|
||||||
|
import torch
|
||||||
|
from transformers import ( # type: ignore
|
||||||
|
AutoModelForVision2Seq,
|
||||||
|
AutoProcessor,
|
||||||
|
BitsAndBytesConfig,
|
||||||
|
)
|
||||||
|
|
||||||
|
device = decide_device(accelerator_options.device)
|
||||||
|
self.device = device
|
||||||
|
|
||||||
|
_log.debug(f"Available device for HuggingFace VLM: {device}")
|
||||||
|
|
||||||
|
repo_cache_folder = vlm_options.repo_id.replace("/", "--")
|
||||||
|
|
||||||
|
# PARAMETERS:
|
||||||
|
if artifacts_path is None:
|
||||||
|
artifacts_path = self.download_models(self.vlm_options.repo_id)
|
||||||
|
elif (artifacts_path / repo_cache_folder).exists():
|
||||||
|
artifacts_path = artifacts_path / repo_cache_folder
|
||||||
|
|
||||||
|
self.param_question = vlm_options.prompt # "Perform Layout Analysis."
|
||||||
|
self.param_quantization_config = BitsAndBytesConfig(
|
||||||
|
load_in_8bit=vlm_options.load_in_8bit, # True,
|
||||||
|
llm_int8_threshold=vlm_options.llm_int8_threshold, # 6.0
|
||||||
|
)
|
||||||
|
self.param_quantized = vlm_options.quantized # False
|
||||||
|
|
||||||
|
self.processor = AutoProcessor.from_pretrained(
|
||||||
|
artifacts_path,
|
||||||
|
# trust_remote_code=True,
|
||||||
|
)
|
||||||
|
if not self.param_quantized:
|
||||||
|
self.vlm_model = AutoModelForVision2Seq.from_pretrained(
|
||||||
|
artifacts_path,
|
||||||
|
device_map=device,
|
||||||
|
torch_dtype=torch.bfloat16,
|
||||||
|
_attn_implementation=(
|
||||||
|
"flash_attention_2"
|
||||||
|
if self.device.startswith("cuda")
|
||||||
|
and accelerator_options.cuda_use_flash_attention2
|
||||||
|
else "eager"
|
||||||
|
),
|
||||||
|
# trust_remote_code=True,
|
||||||
|
) # .to(self.device)
|
||||||
|
|
||||||
|
else:
|
||||||
|
self.vlm_model = AutoModelForVision2Seq.from_pretrained(
|
||||||
|
artifacts_path,
|
||||||
|
device_map=device,
|
||||||
|
torch_dtype="auto",
|
||||||
|
quantization_config=self.param_quantization_config,
|
||||||
|
_attn_implementation=(
|
||||||
|
"flash_attention_2"
|
||||||
|
if self.device.startswith("cuda")
|
||||||
|
and accelerator_options.cuda_use_flash_attention2
|
||||||
|
else "eager"
|
||||||
|
),
|
||||||
|
# trust_remote_code=True,
|
||||||
|
) # .to(self.device)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def download_models(
|
||||||
|
repo_id: str,
|
||||||
|
local_dir: Optional[Path] = None,
|
||||||
|
force: bool = False,
|
||||||
|
progress: bool = False,
|
||||||
|
) -> Path:
|
||||||
|
from huggingface_hub import snapshot_download
|
||||||
|
from huggingface_hub.utils import disable_progress_bars
|
||||||
|
|
||||||
|
if not progress:
|
||||||
|
disable_progress_bars()
|
||||||
|
download_path = snapshot_download(
|
||||||
|
repo_id=repo_id,
|
||||||
|
force_download=force,
|
||||||
|
local_dir=local_dir,
|
||||||
|
# revision="v0.0.1",
|
||||||
|
)
|
||||||
|
|
||||||
|
return Path(download_path)
|
||||||
|
|
||||||
|
def __call__(
|
||||||
|
self, conv_res: ConversionResult, page_batch: Iterable[Page]
|
||||||
|
) -> Iterable[Page]:
|
||||||
|
for page in page_batch:
|
||||||
|
assert page._backend is not None
|
||||||
|
if not page._backend.is_valid():
|
||||||
|
yield page
|
||||||
|
else:
|
||||||
|
with TimeRecorder(conv_res, "vlm"):
|
||||||
|
assert page.size is not None
|
||||||
|
|
||||||
|
hi_res_image = page.get_image(scale=2.0) # 144dpi
|
||||||
|
# hi_res_image = page.get_image(scale=1.0) # 72dpi
|
||||||
|
|
||||||
|
if hi_res_image is not None:
|
||||||
|
im_width, im_height = hi_res_image.size
|
||||||
|
|
||||||
|
# populate page_tags with predicted doc tags
|
||||||
|
page_tags = ""
|
||||||
|
|
||||||
|
if hi_res_image:
|
||||||
|
if hi_res_image.mode != "RGB":
|
||||||
|
hi_res_image = hi_res_image.convert("RGB")
|
||||||
|
|
||||||
|
messages = [
|
||||||
|
{
|
||||||
|
"role": "user",
|
||||||
|
"content": [
|
||||||
|
{
|
||||||
|
"type": "text",
|
||||||
|
"text": "This is a page from a document.",
|
||||||
|
},
|
||||||
|
{"type": "image"},
|
||||||
|
{"type": "text", "text": self.param_question},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
]
|
||||||
|
prompt = self.processor.apply_chat_template(
|
||||||
|
messages, add_generation_prompt=False
|
||||||
|
)
|
||||||
|
inputs = self.processor(
|
||||||
|
text=prompt, images=[hi_res_image], return_tensors="pt"
|
||||||
|
)
|
||||||
|
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
||||||
|
|
||||||
|
start_time = time.time()
|
||||||
|
# Call model to generate:
|
||||||
|
generated_ids = self.vlm_model.generate(
|
||||||
|
**inputs, max_new_tokens=4096, use_cache=True
|
||||||
|
)
|
||||||
|
|
||||||
|
generation_time = time.time() - start_time
|
||||||
|
generated_texts = self.processor.batch_decode(
|
||||||
|
generated_ids[:, inputs["input_ids"].shape[1] :],
|
||||||
|
skip_special_tokens=False,
|
||||||
|
)[0]
|
||||||
|
|
||||||
|
num_tokens = len(generated_ids[0])
|
||||||
|
page_tags = generated_texts
|
||||||
|
|
||||||
|
_log.debug(
|
||||||
|
f"Generated {num_tokens} tokens in time {generation_time:.2f} seconds."
|
||||||
|
)
|
||||||
|
|
||||||
|
# inference_time = time.time() - start_time
|
||||||
|
# tokens_per_second = num_tokens / generation_time
|
||||||
|
# print("")
|
||||||
|
# print(f"Page Inference Time: {inference_time:.2f} seconds")
|
||||||
|
# print(f"Total tokens on page: {num_tokens:.2f}")
|
||||||
|
# print(f"Tokens/sec: {tokens_per_second:.2f}")
|
||||||
|
# print("")
|
||||||
|
page.predictions.vlm_response = VlmPrediction(text=page_tags)
|
||||||
|
|
||||||
|
yield page
|
33
docling/models/hf_vlm_models/pixtral_12b_2409.py
Normal file
33
docling/models/hf_vlm_models/pixtral_12b_2409.py
Normal file
@ -0,0 +1,33 @@
|
|||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from collections.abc import Iterable
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import Page, VlmPrediction
|
||||||
|
from docling.datamodel.document import ConversionResult
|
||||||
|
from docling.datamodel.pipeline_options import (
|
||||||
|
AcceleratorOptions,
|
||||||
|
HuggingFaceVlmOptions,
|
||||||
|
)
|
||||||
|
from docling.models.base_model import BasePageModel
|
||||||
|
from docling.utils.accelerator_utils import decide_device
|
||||||
|
from docling.utils.profiling import TimeRecorder
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class HuggingFaceVlmModel_pixtral_12b_2409(BasePageModel):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
enabled: bool,
|
||||||
|
artifacts_path: Optional[Path],
|
||||||
|
accelerator_options: AcceleratorOptions,
|
||||||
|
vlm_options: HuggingFaceVlmOptions,
|
||||||
|
):
|
||||||
|
self.enabled = enabled
|
||||||
|
|
||||||
|
self.vlm_options = vlm_options
|
||||||
|
|
||||||
|
if self.enabled:
|
||||||
|
import torch
|
@ -24,6 +24,8 @@ from docling.datamodel.settings import settings
|
|||||||
from docling.models.api_vlm_model import ApiVlmModel
|
from docling.models.api_vlm_model import ApiVlmModel
|
||||||
from docling.models.hf_mlx_model import HuggingFaceMlxModel
|
from docling.models.hf_mlx_model import HuggingFaceMlxModel
|
||||||
from docling.models.hf_vlm_model import HuggingFaceVlmModel
|
from docling.models.hf_vlm_model import HuggingFaceVlmModel
|
||||||
|
from docling.models.hf_vlm_models.hf_vlm_model_AutoModelForVision2Seq import HuggingFaceVlmModel_AutoModelForVision2Seq
|
||||||
|
from docling.models.hf_vlm_models.hf_vlm_model_AutoModelForCausalLM import HuggingFaceVlmModel_AutoModelForCausalLM
|
||||||
from docling.pipeline.base_pipeline import PaginatedPipeline
|
from docling.pipeline.base_pipeline import PaginatedPipeline
|
||||||
from docling.utils.profiling import ProfilingScope, TimeRecorder
|
from docling.utils.profiling import ProfilingScope, TimeRecorder
|
||||||
|
|
||||||
@ -77,7 +79,26 @@ class VlmPipeline(PaginatedPipeline):
|
|||||||
vlm_options=vlm_options,
|
vlm_options=vlm_options,
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
|
elif vlm_options.inference_framework == InferenceFramework.TRANSFORMERS_AutoModelForVision2Seq:
|
||||||
|
self.build_pipe = [
|
||||||
|
HuggingFaceVlmModel_AutoModelForVision2Seq(
|
||||||
|
enabled=True, # must be always enabled for this pipeline to make sense.
|
||||||
|
artifacts_path=artifacts_path,
|
||||||
|
accelerator_options=pipeline_options.accelerator_options,
|
||||||
|
vlm_options=vlm_options,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
elif vlm_options.inference_framework == InferenceFramework.TRANSFORMERS_AutoModelForCausalLM:
|
||||||
|
self.build_pipe = [
|
||||||
|
HuggingFaceVlmModel_AutoModelForCausalLM(
|
||||||
|
enabled=True, # must be always enabled for this pipeline to make sense.
|
||||||
|
artifacts_path=artifacts_path,
|
||||||
|
accelerator_options=pipeline_options.accelerator_options,
|
||||||
|
vlm_options=vlm_options,
|
||||||
|
),
|
||||||
|
]
|
||||||
else:
|
else:
|
||||||
|
_log.warning("falling back to HuggingFaceVlmModel (AutoModelForVision2Seq) pipeline")
|
||||||
self.build_pipe = [
|
self.build_pipe = [
|
||||||
HuggingFaceVlmModel(
|
HuggingFaceVlmModel(
|
||||||
enabled=True, # must be always enabled for this pipeline to make sense.
|
enabled=True, # must be always enabled for this pipeline to make sense.
|
||||||
|
@ -7,6 +7,9 @@ from docling_core.types.doc.document import DEFAULT_EXPORT_LABELS
|
|||||||
|
|
||||||
from docling.datamodel.base_models import InputFormat
|
from docling.datamodel.base_models import InputFormat
|
||||||
from docling.datamodel.pipeline_options import (
|
from docling.datamodel.pipeline_options import (
|
||||||
|
InferenceFramework,
|
||||||
|
ResponseFormat,
|
||||||
|
HuggingFaceVlmOptions,
|
||||||
VlmPipelineOptions,
|
VlmPipelineOptions,
|
||||||
smoldocling_vlm_mlx_conversion_options,
|
smoldocling_vlm_mlx_conversion_options,
|
||||||
)
|
)
|
||||||
@ -31,11 +34,30 @@ pipeline_options.force_backend_text = False
|
|||||||
# pipeline_options.vlm_options = smoldocling_vlm_conversion_options
|
# pipeline_options.vlm_options = smoldocling_vlm_conversion_options
|
||||||
|
|
||||||
## Pick a VLM model. Fast Apple Silicon friendly implementation for SmolDocling-256M via MLX
|
## Pick a VLM model. Fast Apple Silicon friendly implementation for SmolDocling-256M via MLX
|
||||||
pipeline_options.vlm_options = smoldocling_vlm_mlx_conversion_options
|
## pipeline_options.vlm_options = smoldocling_vlm_mlx_conversion_options
|
||||||
|
|
||||||
## Alternative VLM models:
|
## Alternative VLM models:
|
||||||
# pipeline_options.vlm_options = granite_vision_vlm_conversion_options
|
# pipeline_options.vlm_options = granite_vision_vlm_conversion_options
|
||||||
|
|
||||||
|
# pixtral_vlm_conversion_options = HuggingFaceVlmOptions(
|
||||||
|
# repo_id="mistralai/Pixtral-12B-Base-2409",
|
||||||
|
# # prompt="OCR the full page to markdown.",
|
||||||
|
# prompt="OCR this image and export it in MarkDown.",
|
||||||
|
# response_format=ResponseFormat.MARKDOWN,
|
||||||
|
# inference_framework=InferenceFramework.TRANSFORMERS,
|
||||||
|
# )
|
||||||
|
|
||||||
|
pixtral_vlm_conversion_options = HuggingFaceVlmOptions(
|
||||||
|
repo_id="microsoft/Phi-4-multimodal-instruct",
|
||||||
|
# prompt="OCR the full page to markdown.",
|
||||||
|
prompt="OCR this image and export it in MarkDown.",
|
||||||
|
response_format=ResponseFormat.MARKDOWN,
|
||||||
|
inference_framework=InferenceFramework.TRANSFORMERS_AutoModelForCausalLM,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
pipeline_options.vlm_options = pixtral_vlm_conversion_options
|
||||||
|
|
||||||
## Set up pipeline for PDF or image inputs
|
## Set up pipeline for PDF or image inputs
|
||||||
converter = DocumentConverter(
|
converter = DocumentConverter(
|
||||||
format_options={
|
format_options={
|
||||||
@ -67,7 +89,7 @@ for source in sources:
|
|||||||
|
|
||||||
for page in res.pages:
|
for page in res.pages:
|
||||||
print("")
|
print("")
|
||||||
print("Predicted page in DOCTAGS:")
|
print(f"Predicted page in {pipeline_options.vlm_options.response_format}:")
|
||||||
print(page.predictions.vlm_response.text)
|
print(page.predictions.vlm_response.text)
|
||||||
|
|
||||||
res.document.save_as_html(
|
res.document.save_as_html(
|
||||||
|
987
poetry.lock
generated
987
poetry.lock
generated
File diff suppressed because it is too large
Load Diff
@ -79,10 +79,7 @@ onnxruntime = [
|
|||||||
{ version = "^1.7.0", optional = true, markers = "python_version >= '3.10'" },
|
{ version = "^1.7.0", optional = true, markers = "python_version >= '3.10'" },
|
||||||
]
|
]
|
||||||
|
|
||||||
transformers = [
|
transformers = "4.42.0"
|
||||||
{ markers = "sys_platform != 'darwin' or platform_machine != 'x86_64'", version = "^4.46.0", optional = true },
|
|
||||||
{ markers = "sys_platform == 'darwin' and platform_machine == 'x86_64'", version = "~4.42.0", optional = true },
|
|
||||||
]
|
|
||||||
accelerate = [
|
accelerate = [
|
||||||
{ markers = "sys_platform != 'darwin' or platform_machine != 'x86_64'", version = "^1.2.1", optional = true },
|
{ markers = "sys_platform != 'darwin' or platform_machine != 'x86_64'", version = "^1.2.1", optional = true },
|
||||||
]
|
]
|
||||||
@ -150,6 +147,11 @@ torchvision = [
|
|||||||
{ markers = "sys_platform == 'darwin' and platform_machine == 'x86_64'", version = "~0.17.2" },
|
{ markers = "sys_platform == 'darwin' and platform_machine == 'x86_64'", version = "~0.17.2" },
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
[tool.poetry.group.lm.dependencies]
|
||||||
|
peft = "^0.15.2"
|
||||||
|
backoff = "^2.2.1"
|
||||||
|
|
||||||
[tool.poetry.extras]
|
[tool.poetry.extras]
|
||||||
tesserocr = ["tesserocr"]
|
tesserocr = ["tesserocr"]
|
||||||
ocrmac = ["ocrmac"]
|
ocrmac = ["ocrmac"]
|
||||||
|
Loading…
Reference in New Issue
Block a user