diff --git a/scripts/query_ai_model.py b/scripts/query_ai_model.py new file mode 100644 index 0000000..dcac18e --- /dev/null +++ b/scripts/query_ai_model.py @@ -0,0 +1,300 @@ +"""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 python_utils import check_env +from data_store.repositories import JobRepository +from data_store.dto import ( + AngleEnum, + ColourEnum, + ContactEnum, + DepthEnum, + DistanceEnum, + LightingEnum, + PointOfViewEnum, +) + +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) + # Calculate accuracy + category_events = {} + for category, initial_prompt in CATEGORY_PROMPT_MAP.items(): + print(f"Calculating accuracy for category: {category}") + accuracy = calculate_model_prompt_accuracy(MODEL, initial_prompt, category) + 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(30): + refined_prompt = refine_prompt(MODEL, prompt_accuracy_map) + print(f"Refined prompt: {refined_prompt}") + accuracy = calculate_model_prompt_accuracy(MODEL, refined_prompt, category) + 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 + print(f"Final prompt accuracy map for category '{category}': {prompt_accuracy_map}") + category_events[category] = prompt_accuracy_map + # save results to json file + with open("prompt_accuracy_results.json", "w", encoding="utf-8") as f: + json.dump(category_events, f, indent=2) + print("Prompt accuracy results saved to 'prompt_accuracy_results.json'.")