from datasets import load_dataset
from PIL import Image

# System message for the assistant
system_message = "You are an expert product description writer for Amazon."

# User prompt that combines the user query and the schema
user_prompt = """Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.
Only return description. The description should be SEO optimized and for a better mobile search experience.

<PRODUCT>
{product}
</PRODUCT>

<CATEGORY>
{category}
</CATEGORY>
"""

# Convert dataset to OAI messages
def format_data(sample):
  return {
      "messages": [
          {
              "role": "system",
              "content": [{"type": "text", "text": system_message}],
          },
          {
              "role": "user",
              "content": [
                  {
                      "type": "text",
                      "text": user_prompt.format(
                          product=sample["Product Name"],
                          category=sample["Category"],
                      ),
                  },
                  {
                      "type": "image",
                      "image": sample["image"],
                  },
              ],
          },
          {
              "role": "assistant",
              "content": [{"type": "text", "text": sample["description"]}],
          },
      ],
  }

def process_vision_info(messages: list[dict]) -> list[Image.Image]:
  image_inputs = []
  # Iterate through each conversation
  for msg in messages:
      # Get content (ensure it's a list)
      content = msg.get("content", [])
      if not isinstance(content, list):
          content = [content]

      # Check each content element for images
      for element in content:
          if isinstance(element, dict) and (
              "image" in element or element.get("type") == "image"
          ):
              # Get the image and convert to RGB
              if "image" in element:
                  image = element["image"]
              else:
                  image = element
              image_inputs.append(image.convert("RGB"))
  return image_inputs

# Load dataset from the hub
dataset = load_dataset("philschmid/amazon-product-descriptions-vlm", split="train")
dataset = dataset.select(range(2))

# Convert dataset to OAI messages
# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
dataset = [format_data(sample) for sample in dataset]


import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig

# Hugging Face model id
model_id = "google/gemma-3-4b-pt" # or `google/gemma-3-12b-pt`, `google/gemma-3-27-pt`

# Check if GPU benefits from bfloat16
if torch.cuda.get_device_capability()[0] < 8:
  raise ValueError("GPU does not support bfloat16, please use a GPU that supports bfloat16.")

# Define model init arguments
model_kwargs = dict(
  attn_implementation="flash_attention_2", # Use "flash_attention_2" when running on Ampere or newer GPU
  torch_dtype=torch.bfloat16, # What torch dtype to use, defaults to auto
  device_map="auto", # Let torch decide how to load the model
)

# Load model and tokenizer
model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
processor = AutoProcessor.from_pretrained("google/gemma-3-4b-it")

from transformers import TrainingArguments

args = TrainingArguments(
  num_train_epochs=1,
  remove_unused_columns=False,
  per_device_train_batch_size=1,
  per_device_eval_batch_size=1,
  bf16=True,
  output_dir="./output",
  eval_strategy="epoch",
  report_to="none",
)

# Create a data collator to encode text and image pairs
def collate_fn(examples):
  texts = []
  images = []
  for example in examples:
      image_inputs = process_vision_info(example["messages"])
      text = processor.apply_chat_template(
          example["messages"], add_generation_prompt=False, tokenize=False
      )
      texts.append(text.strip())
      images.append(image_inputs)

  # Tokenize the texts and process the images
  batch = processor(text=texts, images=images, return_tensors="pt", padding="max_length", max_length=512, truncation=True)

  # The labels are the input_ids, and we mask the padding tokens and image tokens in the loss computation
  labels = batch["input_ids"].clone()

  # Mask image tokens
  image_token_id = [
      processor.tokenizer.convert_tokens_to_ids(
          processor.tokenizer.special_tokens_map["boi_token"]
      )
  ]
  # Mask tokens for not being used in the loss computation
  labels[labels == processor.tokenizer.pad_token_id] = -100
  labels[labels == image_token_id] = -100
  labels[labels == 262144] = -100

  batch["labels"] = labels
  return batch

from transformers import Trainer

trainer = Trainer(
  model=model,
  args=args,
  train_dataset=dataset,
  eval_dataset=dataset,
  processing_class=processor,
  data_collator=collate_fn,
)

# Start training, the model will be automatically saved to the Hub and the output directory
trainer.train()