added use of larger model for refining prompt

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, "
@@ -201,6 +205,13 @@ def calculate_model_prompt_accuracy(
             total_score += score
     # Calculate accuracy
     jobs_evaluated = total_jobs - jobs_skipped
+    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}"
     )
@@ -270,19 +281,51 @@ 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
-        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
+        # Loop to refine prompt based on accuracy
+        for i in range(99):
+            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}")
+            except Exception as e:
+                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:
@@ -290,11 +333,4 @@ if __name__ == "__main__":
                     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'.")
+        mlflow.end_run()