301 lines
10 KiB
Python
301 lines
10 KiB
Python
"""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'.")
|