LLaVA-o1: 시각-언어 모델의 구조적 추론 고도화

"I have an image and a question that I want you to answer. I need you to strictly follow the format with four 
specific sections: SUMMARY, CAPTION, REASONING, and CONCLUSION. It is crucial that you adhere to this structure exactly 
as outlined and that the final answer in the CONCLUSION matches the standard correct answer precisely. 
To explain further: In SUMMARY, briefly explain what steps you'll take to solve the problem. In CAPTION, 
describe the contents of the image, specifically focusing on details relevant to the question. 
In REASONING, outline a step-by-step thought process you would use to solve the problem based on the image. 
In CONCLUSION, give the final answer in a direct format, and it must match the correct answer exactly. 
If it's a multiple choice question, the conclusion should only include the option without repeating what the option is. 
Here's how the format should look: <SUMMARY> [Summarize how you will approach the problem and explain the steps you will 
take to reach the answer.] </SUMMARY> <CAPTION> [Provide a detailed description of the image, particularly emphasizing 
the aspects related to the question.] </CAPTION> <REASONING> [Provide a chain-of-thought, logical explanation of the problem. 
This should outline step-by-step reasoning.] </REASONING> <CONCLUSION> [State the final answer in a clear and direct format. 
It must match the correct answer exactly.] </CONCLUSION> (Do not forget </CONCLUSION>!) Please apply this format meticulously 
to analyze the given image and answer the related question, ensuring that the answer matches the standard one perfectly."

pip install torch==2.4.0 transformers==4.45.0 pillow==10.4.0 datasets==3.1.0 weave==0.51.18 openai==1.55.1


import torch
from transformers import MllamaForConditionalGeneration, AutoProcessor
from PIL import Image
from datasets import load_dataset
import weave

weave.init("llava_single_inference")

@weave.op()
def predict(question: str, image):
    """Run inference with LLaVA on a single sample."""
    # Load model and processor inside the function
    MODEL_ID = "Xkev/Llama-3.2V-11B-cot"
    model = MllamaForConditionalGeneration.from_pretrained(
        MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto"
    )
    processor = AutoProcessor.from_pretrained(MODEL_ID)

    pil_image = Image.open(image) if isinstance(image, str) else image

    # Prepare inputs
    messages = [
        {"role": "user", "content": [
            {"type": "image"},
            {"type": "text", "text": question}
        ]}
    ]
    input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
    inputs = processor(
        pil_image,
        input_text,
        add_special_tokens=False,
        return_tensors="pt"
    ).to(model.device)

    # Generate output
    output = model.generate(
        **inputs,
        max_new_tokens=2048,
        return_dict_in_generate=True,
        output_scores=True
    )

    # Decode output
    output_text = processor.decode(output.sequences[0], skip_special_tokens=True)
    return {"model_output": output_text}


def create_single_sample(dataset_name: str, split: str = "test"):
    """Load a single sample from the dataset."""
    dataset = load_dataset(dataset_name, split=split, cache_dir="./cache")
    example = dataset[0]

    question = example["qa"][0]["query"]
    image = example["image"]

    pil_image = Image.open(image) if isinstance(image, str) else image

    return {
        "question": f"Based on the image, answer the following question: {question}",
        "image": pil_image
    }


if __name__ == "__main__":
    # Load a single sample and perform inference
    sample = create_single_sample("TeeA/ChartQA")
    result = predict(sample["question"], sample["image"])

    # Print the result
    print(f"Question: {sample['question']}")
    print(f"Model Output: {result['model_output']}")


import os
import json
import asyncio
import logging
from datasets import load_dataset
from weave import Evaluation, Model
import weave
from PIL import Image
from io import BytesIO
import base64
import torch
from transformers import MllamaForConditionalGeneration, AutoProcessor
from typing import Dict, Any, ClassVar
from openai import OpenAI
import time

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


weave.init("llava_model_eval")

def get_pil_image(image):
    """Convert various image formats to PIL Image"""
    try:
        if isinstance(image, str):
            return Image.open(image)
        elif isinstance(image, Image.Image):
            return image
        elif isinstance(image, bytes):
            return Image.open(BytesIO(image))
        elif isinstance(image, str) and image.startswith('data:image'):
            base64_data = image.split(',')[1] if ',' in image else image
            image_bytes = base64.b64decode(base64_data)
            return Image.open(BytesIO(image_bytes))
        else:
            raise ValueError(f"Unsupported image type: {type(image)}")
    except Exception as e:
        logger.error(f"Error processing image: {e}")
        raise

def evaluate_chart_answer(question: str, ground_truth: str, predicted_answer: str):
    """Use GPT-4 to evaluate model answers"""
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    
    # # Quick check for exact match in prediction
    # if str(ground_truth).lower() in str(predicted_answer).lower():
    #     return {"correct": 1}
    
    messages = [
        {
            "role": "user",
            "content": f"""You are a judge evaluating answers to chart questions. Return ONLY a JSON object with no additional text.

Question: {question}

Ground Truth Answer: {ground_truth} ###

Predicted Answer: {predicted_answer} 

Compare these answers and return a JSON object in this exact format:
{{"correct": <0 or 1>}}

Rules for scoring:
- Score 1 if the answer is correct (or correct within reasonable error margin)
- Score 1 if the correct answer is contained somewhere in the model response
- Score 1 if the answer is right but the formatting is not exactly the same
- Score 0 if the answer is clearly incorrect"""
        }
    ]
    
    retries = 0
    max_retries = 10
    base_delay = 10

    while retries < max_retries:
        try:
            response = client.chat.completions.create(
                model="gpt-4o-2024-08-06",
                messages=messages,
                max_tokens=1024,
                temperature=0
            )
            
            response_text = response.choices[0].message.content.strip()
            
            try:
                return json.loads(response_text)
            except json.JSONDecodeError:
                import re
                json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
                if json_match:
                    try:
                        return json.loads(json_match.group())
                    except:
                        logger.warning("Failed to parse JSON match")
                logger.warning("Failed to parse response as JSON")
                return {"correct": 0}
        
        except Exception as e:
            logger.error(f"Error during evaluation (attempt {retries + 1}): {e}")
            retries += 1
            if retries < max_retries:
                delay = base_delay * (2 ** (retries - 1))
                logger.info(f"Retrying in {delay} seconds...")
                time.sleep(delay)
            else:
                logger.error("Max retries reached")
                return {"correct": 0}

@weave.op
def chart_qa_scorer(question: str, ground_truth: str, model_output: dict) -> Dict[str, Any]:
    """Weave operation for scoring model outputs"""
    try:
        if not isinstance(model_output, dict) or 'model_output' not in model_output:
            logger.error("Invalid model output format")
            return {'score': 0}
            
        result = evaluate_chart_answer(
            question=question,
            ground_truth=ground_truth,
            predicted_answer=model_output['model_output']
        )
        
        return {'score': result.get('correct', 0)}
        
    except Exception as e:
        logger.error(f"Scoring error: {e}")
        return {'score': 0}

def get_model_prediction(model, question: str, image):
    """Get prediction from model for given question and image"""
    try:
        if not model.model or not model.processor:
            raise ValueError("Model or processor not initialized")
            
        pil_image = get_pil_image(image)
        
        messages = [
            {"role": "user", "content": [
                {"type": "image"},
                {"type": "text", "text": question}
            ]}
        ]
        
        try:
            input_text = model.processor.apply_chat_template(messages, add_generation_prompt=True)
        except Exception as e:
            logger.error(f"Error applying chat template: {e}")
            raise
            
        inputs = model.processor(
            pil_image,
            input_text,
            add_special_tokens=False,
            return_tensors="pt"
        ).to(model.model.device)

        input_tokens = len(inputs.input_ids[0])
        
        output = model.model.generate(
            **inputs, 
            max_new_tokens=2048,
            return_dict_in_generate=True,
            output_scores=True
        )
        
        output_tokens = len(output.sequences[0]) - input_tokens
        output_text = model.processor.decode(output.sequences[0], skip_special_tokens=True)
        
        return {
            'model_output': str(output_text),
            'model': model.__class__.__name__,
            'usage': {
                'prompt_tokens': input_tokens,
                'completion_tokens': output_tokens,
                'total_tokens': input_tokens + output_tokens
            }
        }
        
    except Exception as e:
        logger.error(f"Error during prediction: {e}")
        return {
            'model_output': f"Error during prediction: {str(e)}",
            'model': model.__class__.__name__,
            'usage': {
                'prompt_tokens': 0,
                'completion_tokens': 0,
                'total_tokens': 0
            }
        }

class BaseModel(Model):
    model: ClassVar[MllamaForConditionalGeneration] = None
    processor: ClassVar[AutoProcessor] = None

    @staticmethod
    def load_model():
        logger.info("Loading BaseModel...")
        model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
        try:
            BaseModel.model = MllamaForConditionalGeneration.from_pretrained(
                model_id,
                torch_dtype=torch.bfloat16,
                device_map="auto",
            )
            BaseModel.processor = AutoProcessor.from_pretrained(model_id)
            
            if not BaseModel.model or not BaseModel.processor:
                raise ValueError("Failed to load model or processor")
                
            logger.info("BaseModel loaded successfully")
        except Exception as e:
            logger.error(f"Error loading BaseModel: {e}")
            raise

    @staticmethod
    def clear_model():
        logger.info("Clearing BaseModel from memory...")
        BaseModel.model = None
        BaseModel.processor = None
        torch.cuda.empty_cache()

    @weave.op
    def predict(self, question: str, image):
        if not BaseModel.model or not BaseModel.processor:
            raise ValueError("BaseModel is not loaded.")
            
        out = get_model_prediction(BaseModel, question, image)
        
        try:
            if "assistant" in str(out['model_output']):
                out['model_output'] = str(out['model_output']).split("assistant")[1].strip()
        except Exception as e:
            logger.warning(f"Error splitting output on 'assistant': {e}")
            
        return out

class LlavaModel(Model):
    model: ClassVar[MllamaForConditionalGeneration] = None
    processor: ClassVar[AutoProcessor] = None

    @staticmethod
    def load_model():
        logger.info("Loading LlavaModel...")
        model_id = "Xkev/Llama-3.2V-11B-cot"
        try:
            LlavaModel.model = MllamaForConditionalGeneration.from_pretrained(
                model_id,
                torch_dtype=torch.bfloat16,
                device_map="auto",
            )
            LlavaModel.processor = AutoProcessor.from_pretrained(model_id)
            
            if not LlavaModel.model or not LlavaModel.processor:
                raise ValueError("Failed to load model or processor")
                
            logger.info("LlavaModel loaded successfully")
        except Exception as e:
            logger.error(f"Error loading LlavaModel: {e}")
            raise

    @staticmethod
    def clear_model():
        logger.info("Clearing LlavaModel from memory...")
        LlavaModel.model = None
        LlavaModel.processor = None
        torch.cuda.empty_cache()

    @weave.op
    def predict(self, question: str, image):
        if not LlavaModel.model or not LlavaModel.processor:
            raise ValueError("LlavaModel is not loaded.")
            
        out = get_model_prediction(LlavaModel, question, image)
        
        try:
            if "<CONCLUSION>" in str(out['model_output']):
                out['model_output'] = str(out['model_output']).split("<CONCLUSION>")[1].strip()
        except Exception as e:
            logger.warning(f"Error splitting output on '<CONCLUSION>': {e}")
            
        return out



def create_evaluation_dataset(dataset_name: str, split: str = "test", eval_size: int = 100):
    """
    Create evaluation dataset from the HuggingFace MMVet dataset.
    """
    try:
        # Load and shuffle dataset
        dataset = load_dataset(dataset_name, split=split, cache_dir="./cache").shuffle(seed=42)
        
        # Select the last `eval_size` examples
        start_idx = max(0, len(dataset) - eval_size)
        eval_data = dataset.select(range(0, len(dataset)))

        evaluation_dataset = []
        for example in eval_data:
            question = example["question"]
            ground_truth = example["answer"]
            image = example["image"]

            # Ensure image is in PIL format
            if isinstance(image, str):
                pil_image = Image.open(image)
                if pil_image.mode == "RGBA":
                    pil_image = pil_image.convert("RGB")
            elif isinstance(image, Image.Image):
                pil_image = image
            else:
                raise ValueError(f"Unexpected image type: {type(image)}")

            # Append formatted example to evaluation dataset
            evaluation_dataset.append({
                "question": "Based on the image, answer the following question: " + question,
                "ground_truth": ground_truth,
                "image": pil_image
            })

        if not evaluation_dataset:
            raise ValueError("No examples were processed from the dataset")

        logger.info(f"Created evaluation dataset with {len(evaluation_dataset)} examples")
        return evaluation_dataset

    except Exception as e:
        logger.error(f"Error creating evaluation dataset: {e}")
        raise

async def run_evaluations():
    """Run evaluations on both models"""
    try:
        eval_dataset = create_evaluation_dataset("lmms-lab/MMVet")
        results = {}

        # Evaluate BaseModel
        logger.info("\nEvaluating BaseModel...")
        base_model = BaseModel()
        base_model.load_model()
        evaluation = Evaluation(
            dataset=eval_dataset,
            scorers=[chart_qa_scorer],
            name="BaseModel Evaluation"
        )
        results["base_model"] = await evaluation.evaluate(base_model)
        base_model.clear_model()

        # Evaluate LlavaModel
        logger.info("\nEvaluating LlavaModel...")
        llava_model = LlavaModel()
        llava_model.load_model()
        evaluation = Evaluation(
            dataset=eval_dataset,
            scorers=[chart_qa_scorer],
            name="LlavaModel Evaluation"
        )
        results["llava_model"] = await evaluation.evaluate(llava_model)
        llava_model.clear_model()

        return results

    except Exception as e:
        logger.error(f"Error during evaluation: {e}")
        raise

if __name__ == "__main__":
    try:
        results = asyncio.run(run_evaluations())
        logger.info("Evaluation completed successfully")
        logger.info(f"Results: {results}")
    except Exception as e:
        logger.error(f"Failed to run evaluations: {e}")