atropos/environments/multimodal_dpo/ocr_vqa.py

import base64
import io
import os
import random
import re
import sys
import traceback
from typing import List, Optional, Tuple

from datasets import load_dataset

from atroposlib.envs.base import BaseEnv, BaseEnvConfig, OpenaiConfig, ScoredDataGroup
from atroposlib.type_definitions import GameHistory, Item
from atroposlib.utils.tokenize_for_trainer import tokenize_for_trainer


class OcrVqaEnv(BaseEnv):
    name = "ocr_vqa"
    name_config_cls = BaseEnvConfig

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    async def collect_trajectories(
        self, item: Item
    ) -> Tuple[GameHistory | None, List[Item]]:
        to_score: List[Tuple[GameHistory, str, Optional[str]]] = []
        to_backlog: List[Item] = []

        # Extract question and image from item
        prompt_tuple, gold, base64_image = item
        # The prompt_tuple contains the user prompt as a frozenset
        text_prompt = dict(prompt_tuple[0])["content"]

        # System instruction for answer formatting
        system_msg = {
            "role": "system",
            "content": (
                "You must submit your answer enclosed in <answer> tags, "
                "e.g., <answer>YOUR_ANSWER</answer>"
            ),
        }
        # Multimodal user message with text and image
        user_msg = {
            "role": "user",
            "content": [
                {"type": "text", "text": text_prompt},
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{base64_image}"},
                },
            ],
        }
        messages = [system_msg, user_msg]

        # Call the chat completion endpoint
        chat_completions = await self.server.chat_completion(
            messages=messages,
            n=self.config.group_size,
            max_tokens=512,
            timeout=60,
        )

        # Build trajectories for scoring
        for choice in chat_completions.choices:
            user_hist = {"role": "user", "content": text_prompt}
            assistant_hist = {"role": "assistant", "content": choice.message.content}
            history: GameHistory = (user_hist, assistant_hist)
            to_score.append((history, gold, base64_image))

        return to_score, to_backlog

    async def postprocess_histories(
        self, trajectories: List[GameHistory]
    ) -> ScoredDataGroup:
        # No additional post-processing needed
        pass

    async def evaluate(self, *args, **kwargs):
        # No custom evaluation
        return

    async def setup(self):
        # Load the OCR-VQA dataset
        self.dataset = load_dataset("howard-hou/OCR-VQA")
        self.train = self.dataset["train"]
        self.iter = 0

    async def get_next_item(self) -> Item:
        try:
            entry = self.train[self.iter % len(self.train)]
            self.iter += 1

            # Take the first question and answer
            question = entry["questions"][0]
            answer = entry["answers"][0]
            text_prompt = question
            prompt = tuple(
                [frozenset({"role": "user", "content": text_prompt}.items())]
            )

            # Format the gold answer
            gold_answer = f"<answer>{answer}</answer>"

            # Convert image to base64
            img = entry["image"]
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            img_bytes = buf.getvalue()
            base64_image = base64.b64encode(img_bytes).decode("utf-8")

            return (prompt, gold_answer, base64_image)
        except Exception:
            traceback.print_exc()
            # Fallback example
            fallback_prompt = tuple(
                [
                    frozenset(
                        {"role": "user", "content": "Please solve: 2 + 2 = ?"}.items()
                    )
                ]
            )
            return (fallback_prompt, "<answer>4</answer>", None)

    async def score(self, rollout_group_data) -> Optional[ScoredDataGroup]:
        scores = ScoredDataGroup()
        scores["tokens"] = []
        scores["masks"] = []
        scores["scores"] = []
        scores["images"] = []
        random.shuffle(rollout_group_data)
        for item in rollout_group_data:
            out = tokenize_for_trainer(self.tokenizer, item[0])
            tokens = out["tokens"]
            masks = out["masks"]

            # Extract model and gold answers
            try:
                reply = item[0][-1]["content"]
                m = re.search(r"<answer>\s*(.*?)\s*</answer>", reply, re.IGNORECASE)
                model_answer = m.group(1).strip() if m else reply.strip()

                gold = item[1]
                g = re.search(r"<answer>\s*(.*?)\s*</answer>", gold, re.IGNORECASE)
                gold_answer = g.group(1).strip() if g else gold.strip()

                reward = model_answer.lower() == gold_answer.lower()
            except Exception:
                reward = False

            # Filter out short examples
            if len([i for i in masks if i != -100]) < 10:
                continue

            scores["tokens"].append(tokens)
            scores["masks"].append(masks)
            scores["scores"].append(1.0 if reward else -1.0)
            try:
                scores["images"].append(item[2])
            except Exception:
                scores["images"].append(None)

            if len(scores["tokens"]) >= self.config.group_size:
                break

        return scores

    @classmethod
    def config_init(cls) -> Tuple[BaseEnvConfig, List[OpenaiConfig]]:
        if not os.environ.get("OPENAI_API_KEY"):
            print("ERROR: OPENAI_API_KEY environment variable is not set!")
            sys.exit(1)

        config = BaseEnvConfig(
            wandb_name="ocr_vqa",
            tokenizer_name="gpt2",
            group_size=2,
            use_wandb=False,
            max_num_workers=2,
            rollout_server_url="http://localhost:8000",
            total_steps=1000,
            batch_size=1,
            steps_per_eval=10,
            ensure_scores_are_not_same=False,
        )

        server_configs = [
            OpenaiConfig(
                model_name="gpt-4o",
                base_url=None,
                api_key=os.environ.get("OPENAI_API_KEY"),
                num_requests_for_eval=1,
            ),
        ]

        return config, server_configs


if __name__ == "__main__":
    OcrVqaEnv.cli()