from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig
from PIL import Image
import requests
import torch
from io import BytesIO


model_id = "Ak137/granite-docling-258M-finetuned" 

model = AutoModelForImageTextToText.from_pretrained(
    model_id,
    device_map='cuda',
    attn_implementation="flash_attention_2",
    dtype=torch.bfloat16 
)
processor = AutoProcessor.from_pretrained(model_id)


def generate_text_from_sample(model, processor, sample, max_new_tokens=40):
    prompt = processor.apply_chat_template(sample["prompt"], add_generation_prompt=True)
    inputs = processor(text=prompt, images=sample["images"], return_tensors="pt")
    inputs = inputs.to(model.device)
    
    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
    return (processor.decode(outputs[0][inputs["input_ids"].shape[-1]:]))


user_prompt = "Convert this document to docling format"

url = "https://blogs.loc.gov/law/files/2020/01/sld-misc-image-1.jpg"
response = requests.get(url)
img = Image.open(BytesIO(response.content))
if img.mode != 'RGB':
    img = img.convert('RGB')
    
example = {
            "images": [img],
            "prompt": 
                [
                        {   'role': 'user',
                            'content': 
                                [
                                    {'type': 'image'},
                                    {'text': user_prompt, 'type': 'text'}
                                ],
                         }
                ],
        }


res = generate_text_from_sample(model, processor, sample=example, max_new_tokens=120)
print(res)

Training procedure

This model was trained using SFT using an Nvidia H100 GPU for a total of 4 hours.

Framework versions

TRL: 0.23.1
Transformers: 4.56.2
Pytorch: 2.8.0+cu129
Datasets: 4.1.1
Tokenizers: 0.22.1

Citations

Cite TRL as:

bibtex

@misc{vonwerra2022trl,
	title        = {{TRL: Transformer Reinforcement Learning}},
	author       = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec},
	year         = 2020,
	journal      = {GitHub repository},
	publisher    = {GitHub},
	howpublished = {\url{https://github.com/huggingface/trl}}
}