added code to log to mlflow

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,6 +23,7 @@ from data_store.dto import (
     PointOfViewEnum,
 )
 
+RUN_NAME = f"query-ai-model_{datetime.now().strftime('%Y%m%d-%H%M%S')}"
 MODEL = "llava:13b"
 CATEGORY_PROMPT_MAP = {
     "angle": (
@@ -270,19 +273,49 @@ 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={
+                "model": 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)
+        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(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(MODEL, refined_prompt, category)
+                mlflow.log_metric(f"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 +323,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()