added use of larger model for refining prompt

Reviewed-on: #17

data_store/repositories/annotation_repository.py

@@ -60,17 +60,27 @@ class AnnotationRepository(MinioAdapter, AnnotationRepositoryInterface):
             annotation = Annotation(
                 angle=AngleEnum(data["angle"]) if data["angle"] is not None else None,
                 angle_uncertainty=bool(data["angle_uncertainty"]),
-                colour=ColourEnum(data["colour"]) if data["colour"] is not None else None,
+                colour=ColourEnum(data["colour"])
+                if data["colour"] is not None
+                else None,
                 colour_uncertainty=bool(data["colour_uncertainty"]),
-                contact=ContactEnum(data["contact"]) if data["contact"] is not None else None,
+                contact=ContactEnum(data["contact"])
+                if data["contact"] is not None
+                else None,
                 contact_uncertainty=bool(data["contact_uncertainty"]),
                 depth=DepthEnum(data["depth"]) if data["depth"] is not None else None,
                 depth_uncertainty=bool(data["depth_uncertainty"]),
-                distance=DistanceEnum(data["distance"]) if data["distance"] is not None else None,
+                distance=DistanceEnum(data["distance"])
+                if data["distance"] is not None
+                else None,
                 distance_uncertainty=bool(data["distance_uncertainty"]),
-                lighting=LightingEnum(data["lighting"]) if data["lighting"] is not None else None,
+                lighting=LightingEnum(data["lighting"])
+                if data["lighting"] is not None
+                else None,
                 lighting_uncertainty=bool(data["lighting_uncertainty"]),
-                point_of_view=PointOfViewEnum(data["point_of_view"]) if data["point_of_view"] is not None else None,
+                point_of_view=PointOfViewEnum(data["point_of_view"])
+                if data["point_of_view"] is not None
+                else None,
                 point_of_view_uncertainty=bool(data["point_of_view_uncertainty"]),
             )
         except (KeyError, ValueError) as e:

pyproject.toml

@@ -14,6 +14,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
+    "mlflow>=3.4.0",
     "pandas>=2.3.2",
     "pillow>=11.3.0",
     "pydantic>=2.11.9",
@@ -82,4 +83,5 @@ dev = [
     "ruff>=0.13.0",
     "safety>=3.2.11",
     "testcontainers>=4.13.0",
+    "types-requests>=2.32.4.20250913",
 ]

scripts/query_ai_model.py

@@ -9,6 +9,8 @@ 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 (
@@ -21,7 +23,9 @@ from data_store.dto import (
     PointOfViewEnum,
 )
 
-MODEL = "llava:13b"
+RUN_NAME = f"query-ai-model_{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+IMAGE_MODEL = "llava:13b"
+PROMPT_MODEL = "mistral-small:22b-instruct-2409-q4_K_M"
 CATEGORY_PROMPT_MAP = {
     "angle": (
         "In the context of visual angle in the Kress and van Leeuwen framework, "
@@ -58,7 +62,7 @@ def query_model_with_image(
     prompt: str,
     allowed_answers: list[str],
     image: Image.Image,
-    model: str = "llava:13b"
+    model: str = "llava:13b",
 ) -> str:
     """
     Query an Ollama instance with a prompt, image, and constrained answers.
@@ -77,9 +81,9 @@ def query_model_with_image(
     # Read environment variables
     ollama_endpoint = str(os.getenv("OLLAMA_ENDPOINT"))
     # Handle RGBA images by converting to RGB
-    if image.mode == 'RGBA':
+    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 = 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
@@ -97,14 +101,12 @@ def query_model_with_image(
         "model": model,
         "prompt": prompt_with_options,
         "images": [image_b64],
-        "stream": False
+        "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",
+        f"{ollama_endpoint}/api/generate", json=payload, timeout=60
-        json=payload,
-        timeout=60
     )
     response.raise_for_status()
     result = response.json()
@@ -197,11 +199,22 @@ def calculate_model_prompt_accuracy(
             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}")
+            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}")
+    mlflow.log_metrics(
+        {
+            "total_jobs": total_jobs,
+            "jobs_skipped": jobs_skipped,
+            "jobs_evaluated": jobs_evaluated,
+        }
+    )
+    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
@@ -237,17 +250,11 @@ def refine_prompt(
         "Respond with only the refined prompt."
     )
     # Prepare request payload
-    payload = {
+    payload = {"model": model, "prompt": refinement_instruction, "stream": False}
-        "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",
+        f"{ollama_endpoint}/api/generate", json=payload, timeout=60
-        json=payload,
-        timeout=60
     )
     response.raise_for_status()
     result = response.json()
@@ -274,27 +281,56 @@ if __name__ == "__main__":
     )
     # 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
-    category_events = {}
     for category, initial_prompt in CATEGORY_PROMPT_MAP.items():
+        step = 0
+        run = mlflow.start_run(
+            experiment_id=experiment_id,
+            run_name=RUN_NAME,
+            tags={
+                "image_model": IMAGE_MODEL,
+                "prompt_model": PROMPT_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)
+        accuracy = calculate_model_prompt_accuracy(IMAGE_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
+        # Loop to refine prompt based on accuracy
-        for i in range(30):
+        for i in range(99):
-            refined_prompt = refine_prompt(MODEL, prompt_accuracy_map)
+            step += 1
+            try:
+                # Generate refined prompt
+                refined_prompt = refine_prompt(PROMPT_MODEL, prompt_accuracy_map)
+                mlflow.log_param(f"step_{step}_prompt", refined_prompt)
                 print(f"Refined prompt: {refined_prompt}")
-            accuracy = calculate_model_prompt_accuracy(MODEL, refined_prompt, category)
+            except Exception as e:
-            prompt_accuracy_map[refined_prompt] = accuracy
+                print(f"Error refining prompt: {e}")
+                continue
+            try:
+                # Calculate accuracy for refined prompt
+                accuracy = calculate_model_prompt_accuracy(
+                    IMAGE_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.")
+                print(
+                    f"Achieved 100% accuracy for category '{category}'. Stopping refinement."
+                )
                 break
-        print(f"Final prompt accuracy map for category '{category}': {prompt_accuracy_map}")
+        mlflow.end_run()
-        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'.")