added script to query ai model
This commit is contained in:
Brian Bjarke Jensen
@@ -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'.")
|
||||
Reference in New Issue
Block a user