visual-semiotic-ai-analysis/scripts/query_ai_model.py

"""Script to query an AI model with a prompt and image and print the response."""

import os
import base64
import logging
import json
from io import BytesIO
import requests
import structlog
from dotenv import load_dotenv
from PIL import Image
from datetime import datetime
import mlflow
from python_utils import check_env
from data_store.repositories import JobRepository
from data_store.dto import (
    AngleEnum,
    ColourEnum,
    ContactEnum,
    DepthEnum,
    DistanceEnum,
    LightingEnum,
    PointOfViewEnum,
)

RUN_NAME = f"query-ai-model_{datetime.now().strftime('%Y%m%d-%H%M%S')}"
MODEL = "llava:13b"
CATEGORY_PROMPT_MAP = {
    "angle": (
        "In the context of visual angle in the Kress and van Leeuwen framework, "
        "analyze the image to determine the angle of the subject."
    ),
    "colour": (
        "In the context of colour in the Kress and van Leeuwen framework, "
        "analyze the image to determine the colour saturation of the whole image."
    ),
    "contact": (
        "In the context of contact in the Kress and van Leeuwen framework, "
        "analyze the image to determine the contact between subject and viewer."
    ),
    "depth": (
        "In the context of depth in the Kress and van Leeuwen framework, "
        "analyze the image to determine the depth of the whole image."
    ),
    "distance": (
        "In the context of spatial distance in the Kress and van Leeuwen framework, "
        "analyze the image to determine the spatial distance that the subject is portrayed at."
    ),
    "lighting": (
        "In the context of lighting in the Kress and van Leeuwen framework, "
        "analyze the image to determine the lighting of the whole image."
    ),
    "point_of_view": (
        "In the context of point of view in the Kress and van Leeuwen framework, "
        "analyze the image to determine the orientation of the subject in relation to the viewer."
    ),
}


def query_model_with_image(
    prompt: str,
    allowed_answers: list[str],
    image: Image.Image,
    model: str = "llava:13b",
) -> str:
    """
    Query an Ollama instance with a prompt, image, and constrained answers.

    Args:
        prompt: The text prompt to send
        allowed_answers: List of valid response options
        image: PIL Image to analyze
        model: The Ollama model to use (default: "llava:13b")

    Returns:
        The model's response text
    """
    # Ensure required env vars are set
    check_env("OLLAMA_ENDPOINT")
    # Read environment variables
    ollama_endpoint = str(os.getenv("OLLAMA_ENDPOINT"))
    # Handle RGBA images by converting to RGB
    if image.mode == "RGBA":
        # Convert RGBA to RGB by compositing against white background
        rgb_image = Image.new("RGB", image.size, (255, 255, 255))
        rgb_image.paste(image, mask=image.split()[-1])  # Use alpha as mask
        image = rgb_image
    # Convert PIL Image to base64
    buffer = BytesIO()
    image.save(buffer, format="JPEG")
    image_b64 = base64.b64encode(buffer.getvalue()).decode()
    # Build prompt with allowed answers
    prompt_with_options = (
        f"{prompt}\n"
        f"Please choose one of the following options: {', '.join(allowed_answers)}.\n"
        "Respond with only the chosen option."
    )
    # Prepare request payload
    payload = {
        "model": model,
        "prompt": prompt_with_options,
        "images": [image_b64],
        "stream": False,
    }
    # Make request to Ollama
    # N.B. the generate endpoint does not retain context between calls
    response = requests.post(
        f"{ollama_endpoint}/api/generate", json=payload, timeout=60
    )
    response.raise_for_status()
    result = response.json()

    return str(result.get("response", "").strip())


def score_response(response: str, expected_answer: str) -> int:
    """
    Score the model's response based on allowed answers.

    Args:
        response: The model's response text
        expected_answer: The expected answer

    Returns:
        1 if the response is in allowed answers, else 0
    """
    # Score response
    return 1 if response.lower() == expected_answer.lower() else 0


def allowed_answers_for_category(
    category: str,
) -> list[str]:
    """
    Get allowed answers based on the category.

    Args:
        category: The category to determine allowed answers

    Returns:
        List of allowed answers
    """
    match category.lower():
        case "angle":
            return [e.value for e in AngleEnum]
        case "colour":
            return [e.value for e in ColourEnum]
        case "contact":
            return [e.value for e in ContactEnum]
        case "depth":
            return [e.value for e in DepthEnum]
        case "distance":
            return [e.value for e in DistanceEnum]
        case "lighting":
            return [e.value for e in LightingEnum]
        case "point_of_view":
            return [e.value for e in PointOfViewEnum]
        case _:
            raise ValueError(f"Invalid category: {category}")


def calculate_model_prompt_accuracy(
    model: str,
    prompt: str,
    category: str,
) -> float:
    """
    Calculate the accuracy of the prompt against the expected answers for the given category.
    """
    # Determine allowed answers based on category
    allowed_answers = allowed_answers_for_category(category)
    # Score each job in the repository)
    jobs_skipped = 0
    total_score = 0
    with JobRepository() as job_repo:
        job_ids = job_repo.list_all()
        if not job_ids:
            raise ValueError("No jobs found in JobRepository.")
        total_jobs = len(job_ids)
        # Get all responses for the prompt
        for job_id in job_ids:
            job = job_repo.get(job_id)
            if not job:
                jobs_skipped += 1
                continue
            # skip job if uncertainty flag is set for the category
            if getattr(job.annotation, f"{category.lower()}_uncertainty"):
                print(f"Job {job_id}: Skipping due to uncertainty flag.")
                jobs_skipped += 1
                continue
            response = query_model_with_image(
                model=model,
                prompt=prompt,
                allowed_answers=allowed_answers,
                image=job.image,
            )
            # Get expected answer from job annotation
            expected_answer = getattr(job.annotation, category.lower())
            # Score the response
            score = score_response(response, expected_answer)
            print(
                f"Job {job_id}: Response '{response}' - Expected answer '{expected_answer}' - Score: {score}"
            )
            total_score += score
    # Calculate accuracy
    jobs_evaluated = total_jobs - jobs_skipped
    print(
        f"Total Jobs: {total_jobs}, Jobs Skipped: {jobs_skipped}, Jobs Evaluated: {jobs_evaluated}"
    )
    if jobs_evaluated == 0:
        raise ValueError("No jobs were evaluated. All jobs were skipped.")
    accuracy = total_score / jobs_evaluated
    return accuracy


def refine_prompt(
    model: str,
    prompt_accuracy_map: dict[str, float],
) -> str:
    """
    Refine the prompt to improve accuracy.

    Args:
        model: The Ollama model to use
        prompt_accuracy_map: A mapping of prompts to their accuracies

    Returns:
        The refined prompt text
    """
    # Ensure required env vars are set
    check_env("OLLAMA_ENDPOINT")
    # Read environment variables
    ollama_endpoint = str(os.getenv("OLLAMA_ENDPOINT"))
    # Convert prompt accuracy map to json string for better readability
    prompt_accuracy_json = json.dumps(prompt_accuracy_map, indent=2)
    # Build refinement instruction
    refinement_instruction = (
        "Given the following prompts and accuracies:\n"
        f"{prompt_accuracy_json}\n"
        "Please refine the prompt to improve the accuracy.\n"
        "The refined prompt should be clear, concise, and focused on obtaining accurate responses.\n"
        "Respond with only the refined prompt."
    )
    # Prepare request payload
    payload = {"model": model, "prompt": refinement_instruction, "stream": False}
    # Make request to Ollama
    # N.B. the generate endpoint does not retain context between calls
    response = requests.post(
        f"{ollama_endpoint}/api/generate", json=payload, timeout=60
    )
    response.raise_for_status()
    result = response.json()

    return str(result.get("response", "").strip())


if __name__ == "__main__":
    load_dotenv()
    # Ensure required env vars are set
    check_env("OLLAMA_ENDPOINT")
    # Configure structlog
    structlog.configure(
        processors=[
            structlog.stdlib.filter_by_level,
            structlog.stdlib.add_logger_name,
            structlog.stdlib.add_log_level,
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.processors.JSONRenderer(),
        ],
        logger_factory=structlog.stdlib.LoggerFactory(),
        wrapper_class=structlog.stdlib.BoundLogger,
        cache_logger_on_first_use=True,
    )
    # Set up basic logging configuration
    logging.basicConfig(level=logging.ERROR)
    # Connect to MLFlow
    mlflow.set_tracking_uri("http://10.0.0.2:5000")
    experiment = mlflow.get_experiment_by_name("visual-semiotic-ai-analysis")
    experiment_id = experiment.experiment_id if experiment else None
    # Calculate accuracy
    for category, initial_prompt in CATEGORY_PROMPT_MAP.items():
        step = 0
        run = mlflow.start_run(
            experiment_id=experiment_id,
            run_name=RUN_NAME,
            tags={
                "model": MODEL,
                "category": category,
            },
        )
        mlflow.log_param(f"step_{step}_prompt", initial_prompt)
        print(f"Calculating accuracy for category: {category}")
        accuracy = calculate_model_prompt_accuracy(MODEL, initial_prompt, category)
        mlflow.log_metric("accuracy", accuracy, step=step)
        # Log to MLFlow
        prompt_accuracy_map = {initial_prompt: accuracy}
        print(f"Prompt accuracy for category '{category}': {accuracy:.2%}")
        # Loop to refine prompt based on accuracy
        for i in range(99):
            step += 1
            try:
                # Generate refined prompt
                refined_prompt = refine_prompt(MODEL, prompt_accuracy_map)
                mlflow.log_param(f"step_{step}_prompt", refined_prompt)
                print(f"Refined prompt: {refined_prompt}")
            except Exception as e:
                print(f"Error refining prompt: {e}")
                continue
            try:
                # Calculate accuracy for refined prompt
                accuracy = calculate_model_prompt_accuracy(
                    MODEL, refined_prompt, category
                )
                mlflow.log_metric("accuracy", accuracy, step=step)
            except Exception as e:
                print(f"Error calculating accuracy: {e}")
                continue
            # Update prompt accuracy map
            # prompt_accuracy_map[refined_prompt] = accuracy
            print(f"Prompt accuracy for category '{category}': {accuracy:.2%}")
            # Stop if accuracy is 100%
            if accuracy == 1.0:
                print(
                    f"Achieved 100% accuracy for category '{category}'. Stopping refinement."
                )
                break
        mlflow.end_run()