使用视觉语言模型从图像或文档进行结构化生成

我们将使用 HuggingFaceTB 的 SmolVLM-Instruct 模型从文档中提取结构化信息。我们将使用 Hugging Face Transformers 库和 Outlines 库运行 VLM，Outlines 库有助于基于限制令牌采样概率的结构化生成。

此方法基于 Outlines 教程。

依赖项和导入

首先，让我们安装必要的库。

%pip install accelerate outlines transformers torch flash-attn datasets sentencepiece

接下来，导入必要的库。

import outlines
import torch

from datasets import load_dataset
from outlines.models.transformers_vision import transformers_vision
from transformers import AutoModelForImageTextToText, AutoProcessor
from pydantic import BaseModel

初始化模型

我们将从 HuggingFaceTB/SmolVLM-Instruct 初始化模型。Outlines 要求我们传入模型类和处理器类，因此我们将通过创建一个返回这些的函数来使此示例更通用。或者，您可以查看 Hub repo 文件中的模型和分词器配置，并直接导入这些类。

model_name = "HuggingFaceTB/SmolVLM-Instruct"


def get_model_and_processor_class(model_name: str):
    model = AutoModelForImageTextToText.from_pretrained(model_name)
    processor = AutoProcessor.from_pretrained(model_name)
    classes = model.__class__, processor.__class__
    del model, processor
    return classes


model_class, processor_class = get_model_and_processor_class(model_name)

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

model = transformers_vision(
    model_name,
    model_class=model_class,
    device=device,
    model_kwargs={"torch_dtype": torch.bfloat16, "device_map": "auto"},
    processor_kwargs={"device": device},
    processor_class=processor_class,
)

结构化生成

现在，我们将定义一个函数，该函数将定义我们模型的输出结构。我们将使用 openbmb/RLAIF-V-Dataset，其中包含一组图像以及问题及其选择和拒绝的响应。这是一个不错的数据集，但我们希望在图像之上创建额外的图像到文本数据，以获取我们自己的结构化数据集，并可能在此基础上微调我们的模型。我们将使用模型为图像生成字幕、问题和简单的质量标签。

class ImageData(BaseModel):
    quality: str
    description: str
    question: str


structured_generator = outlines.generate.json(model, ImageData)

现在，让我们想出一个提取提示。

prompt = """
You are an image analysis assisant.

Provide a quality tag, a description and a question.

The quality can either be "good", "okay" or "bad".
The question should be concise and objective.

Return your response as a valid JSON object.
""".strip()

让我们加载图像数据集。

dataset = load_dataset("openbmb/RLAIF-V-Dataset", split="train[:10]")
dataset

现在，让我们定义一个函数，该函数将从图像中提取结构化信息。我们将使用 `apply_chat_template` 方法格式化提示，然后将其连同图像一起传递给模型。

def extract(row):
    messages = [
        {
            "role": "user",
            "content": [{"type": "image"}, {"type": "text", "text": prompt}],
        },
    ]

    formatted_prompt = model.processor.apply_chat_template(messages, add_generation_prompt=True)

    result = structured_generator(formatted_prompt, [row["image"]])
    row["synthetic_question"] = result.question
    row["synthetic_description"] = result.description
    row["synthetic_quality"] = result.quality
    return row


dataset = dataset.map(lambda x: extract(x))
dataset

现在，让我们将新数据集推送到 Hub。

dataset.push_to_hub(
    "davidberenstein1957/structured-generation-information-extraction-vlms-openbmb-RLAIF-V-Dataset", split="train"
)

结果并不完美，但它们是继续探索不同模型和提示的良好起点！

结论

我们已经看到了如何使用视觉语言模型从文档中提取结构化信息。我们可以使用类似的提取方法从文档中提取结构化信息，例如使用 `pdf2image` 将文档转换为图像，并对每页 PDF 图像进行信息提取。

pdf_path = "path/to/your/pdf/file.pdf"
pages = convert_from_path(pdf_path)
for page in pages:
    extract_objects = extract_objects(page, prompt)

后续步骤

请查看 Outlines 库以获取更多关于如何使用它的信息。探索不同的方法和参数。
使用您自己的模型探索您自己的用例中的提取。
使用不同的方法从文档中提取结构化信息。

< > 在 GitHub 上更新

开源 AI 食谱

使用视觉语言模型从图像或文档进行结构化生成

依赖项和导入

初始化模型

结构化生成

结论

后续步骤