From 5faf3ef2a75ad8ce9403ba502ff13f42eef20f2c Mon Sep 17 00:00:00 2001
From: brian <schnitzelen@gmail.com>
Date: Fri, 26 Sep 2025 19:08:56 +0000
Subject: [PATCH] added use of larger model for refining prompt

---
 scripts/query_ai_model.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/scripts/query_ai_model.py b/scripts/query_ai_model.py
index 5a3424f..7503f3d 100644
--- a/scripts/query_ai_model.py
+++ b/scripts/query_ai_model.py
@@ -24,7 +24,8 @@ from data_store.dto import (
 )
 
 RUN_NAME = f"query-ai-model_{datetime.now().strftime('%Y%m%d-%H%M%S')}"
-MODEL = "llava:13b"
+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, "
@@ -204,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}"
     )
@@ -284,13 +292,14 @@ if __name__ == "__main__":
             experiment_id=experiment_id,
             run_name=RUN_NAME,
             tags={
-                "model": MODEL,
+                "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}
@@ -300,7 +309,7 @@ if __name__ == "__main__":
             step += 1
             try:
                 # Generate refined prompt
-                refined_prompt = refine_prompt(MODEL, prompt_accuracy_map)
+                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:
@@ -309,7 +318,7 @@ if __name__ == "__main__":
             try:
                 # Calculate accuracy for refined prompt
                 accuracy = calculate_model_prompt_accuracy(
-                    MODEL, refined_prompt, category
+                    IMAGE_MODEL, refined_prompt, category
                 )
                 mlflow.log_metric("accuracy", accuracy, step=step)
             except Exception as e: