Working with Pixtral Large for visual chart understanding

pip install -U mistralai weave pillow datasets openai anthropic

import base64
from io import BytesIO
from PIL import Image
from datasets import load_dataset
from mistralai import Mistral
import weave

# Initialize Weave
weave.init("pixtral_inference")

# Set up the Pixtral API client
API_KEY = "your api key"  # Replace with your API key
client = Mistral(api_key=API_KEY)

# Helper function to encode an image to base64
def encode_image(image_obj):
    if isinstance(image_obj, Image.Image):  # Check if it's already a PIL Image
        img = image_obj
    else:  # Otherwise, try opening it as a path
        img = Image.open(image_obj)
    buffered = BytesIO()
    img.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

# Define a function to perform inference
@weave.op
def run_inf(image_obj, prompt):
    # Encode the image to base64
    image_base64 = encode_image(image_obj)

    # Prepare input for the Pixtral API
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}}
            ]
        }
    ]

    # Perform inference
    response = client.chat.complete(
        model="pixtral-large-latest",
        messages=messages,
        max_tokens=300
    )
    
    # Return the model's output
    return response.choices[0].message.content

# Example usage with the dataset
if __name__ == "__main__":
    # Load the ChartQA dataset
    dataset = load_dataset("TeeA/ChartQA", split="train")

    # Select one sample for inference
    sample_data = dataset[0]  # Get the first data point

    # Extract the image and prompt
    image_obj = sample_data["image"]  # This should be a PIL Image or file path
    prompt = sample_data["qa"][0]["query"]  # Get the query for the first QA pair

    # Perform inference
    result = run_inf(image_obj, prompt)
    
    # Print the result
    print(f"Question: {prompt}")
    print(f"Model Prediction: {result}")

import os
import base64
from io import BytesIO 
from PIL import Image

class EzModel:
    def __init__(self, model_name, max_tokens=1024, temperature=0.0):
        self.model_name = model_name.lower()
        self.api_key = os.getenv(f"{model_name.upper()}_API_KEY")
        self.max_tokens = max_tokens
        self.temperature = temperature
        
        if not self.api_key:
            raise ValueError(f"API key for {model_name} not found in environment variables.")
        self.client = self._initialize_client()

    def _initialize_client(self):
        if self.model_name == "claude":
            from anthropic import Anthropic
            return Anthropic(api_key=self.api_key)
        elif self.model_name == "gpt":
            from openai import OpenAI
            return OpenAI(api_key=self.api_key)
        elif self.model_name == "mistral":
            from mistralai import Mistral
            return Mistral(api_key=self.api_key)
        else:
            raise ValueError(f"Unsupported model name: {self.model_name}")
    
    def _encode_image(self, image):
        """Convert PIL image to base64 with correct format per API"""
        if image.mode == "RGBA":
            image = image.convert("RGB")
        
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        return base64.b64encode(buffered.getvalue()).decode("utf-8")

    def __call__(self, prompt, pillage=None, base64Img=None, max_tokens=None, temperature=None):
        # Handle image input
        image_base64 = None
        if pillage:
            if isinstance(pillage, Image.Image):
                image_base64 = self._encode_image(pillage)
            else:
                raise ValueError("pillage must be a PIL Image object")
        elif base64Img:
            image_base64 = base64Img

        max_tokens = max_tokens if max_tokens is not None else self.max_tokens
        temperature = temperature if temperature is not None else self.temperature

        if self.model_name == "claude":
            return self._infer_claude(image_base64, prompt, max_tokens, temperature)
        elif self.model_name == "gpt":
            return self._infer_gpt(image_base64, prompt, max_tokens, temperature)
        elif self.model_name == "mistral":
            return self._infer_mistral(image_base64, prompt, max_tokens, temperature)
        else:
            raise ValueError(f"Unsupported model: {self.model_name}")

    def _infer_claude(self, image_base64, prompt, max_tokens, temperature):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt
                    }
                ]
            }
        ]
        if image_base64:
            messages[0]["content"].insert(0, {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/jpeg",
                    "data": image_base64
                }
            })
            
        return self.client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=max_tokens,
            temperature=temperature,
            messages=messages
        ).content[0].text

    def _infer_gpt(self, image_base64, prompt, max_tokens, temperature):
        if image_base64:
            content = [
                {
                    "type": "text",
                    "text": prompt
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{image_base64}"
                    }
                }
            ]
        else:
            content = prompt
            
        return self.client.chat.completions.create(
            model="gpt-4o-2024-08-06",
            messages=[{"role": "user", "content": content}],
            max_tokens=max_tokens,
            temperature=temperature
        ).choices[0].message.content

    def _infer_mistral(self, image_base64, prompt, max_tokens, temperature):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt
                    }
                ]
            }
        ]
        if image_base64:
            messages[0]["content"].append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{image_base64}"
                }
            })
            
        return self.client.chat.complete(
            model="pixtral-large-latest",
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature
        ).choices[0].message.content

import os
import json
import asyncio
from datasets import load_dataset
from ez_model import EzModel
from weave import Evaluation, Model
import weave
from PIL import Image
from io import BytesIO
import base64

# Initialize Weave
weave.init("pixtral_vs_others_eval")

def get_pil_image(image):
    """Convert input to PIL Image"""
    if isinstance(image, str):  # File path
        return Image.open(image)
    elif isinstance(image, Image.Image):  # Already PIL image
        return image
    elif isinstance(image, bytes):  # base64 decoded bytes
        return Image.open(BytesIO(image))
    elif isinstance(image, str) and image.startswith('data:image'):  # base64 string
        # Strip data URL prefix if present
        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("Unsupported image type")

# Instantiate shared model instances
model_instances = {
    "claude": EzModel("claude"),
    "gpt": EzModel("gpt"),
    "mistral": EzModel("mistral"),
}

class ClaudeModel(Model):
    @weave.op()
    def predict(question: str, image):
        # Hardcoded model instance for "claude"
        model = model_instances["claude"]
        # Convert image to PIL and pass as pillage
        pil_image = get_pil_image(image)
        response = model(prompt=question, pillage=pil_image)
        return {"model_output": response}

class GPTModel(Model):
    @weave.op()
    def predict(question: str, image):
        # Hardcoded model instance for "gpt"
        model = model_instances["gpt"]
        # Convert image to PIL and pass as pillage
        pil_image = get_pil_image(image)
        response = model(prompt=question, pillage=pil_image)
        return {"model_output": response}

class MistralModel(Model):
    @weave.op()
    def predict(question: str, image):
        # Hardcoded model instance for "mistral"
        model = model_instances["mistral"]
        # Convert image to PIL and pass as pillage
        pil_image = get_pil_image(image)
        response = model(prompt=question, pillage=pil_image)
        return {"model_output": response}
    

@weave.op()
def llm_judge_scorer(ground_truth: str, model_output: dict) -> dict:
    """Query multiple LLMs to judge correctness and use Claude to determine the majority vote."""
    if not model_output or "model_output" not in model_output:
        return {"llm_judge_score": 0.0}

    # Gather individual judgments from all models
    judgments = {}
    for judge_name, judge_model in model_instances.items():
        prompt = (
            f"You are an LLM judge evaluating correctness on ChartQA.\n"
            f"Ground Truth: {ground_truth}\n"
            f"Predicted Response: {model_output['model_output']}\n\n"
            "Does the predicted response contain the ground truth answer? "
            "The answer does not need to be verbatim but must contain the correct answer. "
            "Respond with 'true' or 'false'."
        )
        response = judge_model(prompt=prompt)
        judgments[judge_name] = "true" in response.lower().strip()

    
    majority_prompt = (
        f"You are tasked with determining the majority vote from multiple LLM judgments.\n"
        f"The judgments from the LLMs are as follows:\n{json.dumps(judgments)}\n\n"
        f"Determine the majority vote based on the judgments. "
        f"If there is a tie, the decision is 'false'. "
        f"Respond with 'true' or 'false'."
    )
    final_decision = model_instances["claude"](prompt=majority_prompt)


    return {"llm_judge_score": 1.0 if "true" in final_decision.lower().strip() else 0.0}


def create_evaluation_dataset(dataset_name: str, split: str = "test", eval_size: int = 3):
    """Prepare evaluation dataset from the ChartQA dataset."""
    dataset = load_dataset(dataset_name, split=split, cache_dir="./cache").shuffle(seed=42)
    
    # Get last eval_size examples
    start_idx = max(0, len(dataset) - eval_size)
    eval_data = dataset.select(range(start_idx, len(dataset)))

    evaluation_dataset = []
    for example in eval_data:
        question = example["qa"][0]["query"]
        ground_truth = example["qa"][0]["label"]
        image = example["image"]

        # Convert to PIL Image if needed
        if isinstance(image, str):  # If it's a file path
            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)}")

        evaluation_dataset.append({
            "question": "Based on the image, answer the following question: " + question,
            "ground_truth": ground_truth,
            "image": pil_image
        })

    # Take first 3 examples and ensure we have data
    if not evaluation_dataset:
        raise ValueError("No examples were processed from the dataset")
        
    return evaluation_dataset

async def run_evaluations():
    """Run evaluations for all models."""
    eval_dataset = create_evaluation_dataset("TeeA/ChartQA")

    # Initialize models
    models = {
        "claude": ClaudeModel(),
        "gpt": GPTModel(),
        "mistral": MistralModel(),
    }

    # Define scorers
    scorers = [llm_judge_scorer]

    # Run evaluations
    results = {}
    for model_name, model in models.items():
        print(f"\nEvaluating {model_name}...")
        evaluation = Evaluation(
            dataset=eval_dataset,
            scorers=scorers,
            name=model_name + " Evaluation"
        )
        results[model_name] = await evaluation.evaluate(model.predict)

    # Print and save results
    print("\nEvaluation Results:")
    for model_name, result in results.items():
        print(f"\n{model_name} Results:")
        print(json.dumps(result, indent=2))

    output_file = "image_model_evaluation_results.json"
    with open(output_file, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {output_file}")

    return results

if __name__ == "__main__":
    asyncio.run(run_evaluations())