fluffy-koala

Model Card for Model ID

In order to run this notebook you need to first create an .env file in the root directory with the following content.

HF_TOKEN=READ_WRITE_HF_TOKEN
TOKENIZERS_PARALLELISM=false
WANDB_API_KEY=WANDB_TOKEN_FOR_LOGGING

HF_TOKEN is needed for pushing the model to HF. WANDB_API_KEY is needed for evaluating the accuracy of fine tuning.

Before running the fine tuning you need to generate fake data. To do that install the requirements.txt dependencies. Since the file was created on MacOS, the installation might fail on Linux system. Make sure to install faker (pip install faker) for the script to run properly.

python dataset_generator.py

Once these steps have been completed, you can run the fine tuning notebook. Notebook was tested with Python 3.10.

Running inference

from transformers import PaliGemmaForConditionalGeneration, AutoProcessor, BitsAndBytesConfig
from PIL import Image
import torch

image = Image.open('path/to/image')

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
model = PaliGemmaForConditionalGeneration.from_pretrained("theDisco/fluffy-koala", quantization_config=bnb_config)
processor = AutoProcessor.from_pretrained("google/paligemma-3b-pt-448")

inputs = processor(text="extract JSON.", images=image, return_tensors="pt")
generated_ids = model.generate(**inputs, max_new_tokens=512)

image_token_index = model.config.image_token_index
num_image_tokens = len(generated_ids[generated_ids == image_token_index])
num_text_tokens = len(processor.tokenizer.encode("extract JSON."))
num_prompt_tokens = num_image_tokens + num_text_tokens + 2
generated_text = processor.batch_decode(generated_ids[:, num_prompt_tokens:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

print(generated_text)