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

我们将使用 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 仓库文件中的模型和分词器配置，并直接导入这些类。

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 食谱

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

依赖项和导入

初始化我们的模型

结构化生成

结论

下一步