"""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'.")