palimpzest 0.7.7__py3-none-any.whl → 0.7.9__py3-none-any.whl

palimpzest/constants.py

@@ -10,21 +10,69 @@ class Model(str, Enum):
     which requires invoking an LLM. It does NOT specify whether the model need be executed
     remotely or locally (if applicable).
     """
-    # LLAMA3 = "meta-llama/Llama-3-8b-chat-hf"
-    LLAMA3 = "meta-llama/Llama-3.3-70B-Instruct-Turbo"
-    LLAMA3_V = "meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo"
+    LLAMA3_2_3B = "meta-llama/Llama-3.2-3B-Instruct-Turbo"
+    LLAMA3_1_8B = "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo"
+    LLAMA3_3_70B = "meta-llama/Llama-3.3-70B-Instruct-Turbo"
+    LLAMA3_2_90B_V = "meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo"
     MIXTRAL = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-    DEEPSEEK = "deepseek-ai/DeepSeek-V3"
+    DEEPSEEK_V3 = "deepseek-ai/DeepSeek-V3"
+    DEEPSEEK_R1_DISTILL_QWEN_1_5B = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
     GPT_4o = "gpt-4o-2024-08-06"
-    GPT_4o_V = "gpt-4o-2024-08-06"
     GPT_4o_MINI = "gpt-4o-mini-2024-07-18"
-    GPT_4o_MINI_V = "gpt-4o-mini-2024-07-18"
     TEXT_EMBEDDING_3_SMALL = "text-embedding-3-small"
     CLIP_VIT_B_32 = "clip-ViT-B-32"
+    # o1 = "o1-2024-12-17"
 
     def __repr__(self):
         return f"{self.name}"
 
+    def is_deepseek_model(self):
+        return "deepseek" in self.value.lower()
+
+    def is_llama_model(self):
+        return "llama" in self.value.lower()
+
+    def is_mixtral_model(self):
+        return "mixtral" in self.value.lower()
+
+    def is_clip_model(self):
+        return "clip" in self.value.lower()
+
+    def is_together_model(self):
+        is_llama_model = self.is_llama_model()
+        is_mixtral_model = self.is_mixtral_model()
+        is_deepseek_model = self.is_deepseek_model()
+        is_clip_model = self.is_clip_model()
+        return is_llama_model or is_mixtral_model or is_deepseek_model or is_clip_model
+
+    def is_gpt_4o_model(self):
+        return "gpt-4o" in self.value.lower()
+
+    def is_o1_model(self):
+        return "o1" in self.value.lower()
+
+    def is_text_embedding_model(self):
+        return "text-embedding" in self.value.lower()
+
+    def is_openai_model(self):
+        is_gpt4_model = self.is_gpt_4o_model()
+        is_o1_model = self.is_o1_model()
+        is_text_embedding_model = self.is_text_embedding_model()
+        return is_gpt4_model or is_o1_model or is_text_embedding_model
+
+    def is_vision_model(self):
+        vision_models = [
+            "meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo",
+            "gpt-4o-2024-08-06",
+            "gpt-4o-mini-2024-07-18",
+            "o1-2024-12-17",
+        ]
+        return self.value in vision_models
+
+    def is_embedding_model(self):
+        is_clip_model = self.is_clip_model()
+        is_text_embedding_model = self.is_text_embedding_model()
+        return is_clip_model or is_text_embedding_model
 
 class APIClient(str, Enum):
     """
@@ -194,23 +242,10 @@ LOG_LLM_OUTPUT = False
 
 
 #### MODEL PERFORMANCE & COST METRICS ####
-# I've looked across models and grouped knowledge into commonly used categories:
-# - Agg. Benchmark (we only use MMLU for this)
-# - Commonsense Reasoning
-# - World Knowledge
-# - Reading Comprehension
-# - Code
-# - Math
-#
-# We don't have global overlap on the World Knowledge and/or Reading Comprehension
-# datasets. Thus, we include these categories results where we have them, but they
-# are completely omitted for now.
-#
-# Within each category only certain models have overlapping results on the same
-# individual datasets; in order to have consistent evaluations I have computed
-# the average result for each category using only the shared sets of datasets within
-# that category. All datasets for which we have results will be shown but commented
-# with ###; datasets which are used in our category averages will have a ^.
+# Overall model quality is computed using MMLU-Pro; multi-modal models currently use the same score for vision
+# - in the future we should split quality for vision vs. multi-modal vs. text
+# - code quality was computed using HumanEval, but that benchmark is too easy and should be replaced.
+# - https://huggingface.co/spaces/TIGER-Lab/MMLU-Pro
 #
 # Cost is presented in terms of USD / token for input tokens and USD / token for
 # generated tokens.
@@ -220,17 +255,28 @@ LOG_LLM_OUTPUT = False
 # values more precisely:
 # - https://artificialanalysis.ai/models/llama-3-1-instruct-8b
 #
-# LLAMA3_8B_MODEL_CARD = {
-#     ##### Cost in USD #####
-#     "usd_per_input_token": 0.18 / 1E6,
-#     "usd_per_output_token": 0.18 / 1E6,
-#     ##### Time #####
-#     "seconds_per_output_token": 0.0061,
-#     ##### Agg. Benchmark #####
-#     "overall": 71.0,
-#     ##### Code #####
-#     "code": 64.0,
-# }
+LLAMA3_2_3B_INSTRUCT_MODEL_CARD = {
+    ##### Cost in USD #####
+    "usd_per_input_token": 0.06 / 1e6,
+    "usd_per_output_token": 0.06 / 1e6,
+    ##### Time #####
+    "seconds_per_output_token": 0.0064,
+    ##### Agg. Benchmark #####
+    "overall": 36.50, # https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct/discussions/13
+    ##### Code #####
+    "code": 0.0,
+}
+LLAMA3_1_8B_INSTRUCT_MODEL_CARD = {
+    ##### Cost in USD #####
+    "usd_per_input_token": 0.18 / 1e6,
+    "usd_per_output_token": 0.18 / 1e6,
+    ##### Time #####
+    "seconds_per_output_token": 0.0059,
+    ##### Agg. Benchmark #####
+    "overall": 44.25,
+    ##### Code #####
+    "code": 72.6,
+}
 LLAMA3_3_70B_INSTRUCT_MODEL_CARD = {
     ##### Cost in USD #####
     "usd_per_input_token": 0.88 / 1e6,
@@ -238,19 +284,10 @@ LLAMA3_3_70B_INSTRUCT_MODEL_CARD = {
     ##### Time #####
     "seconds_per_output_token": 0.0139,
     ##### Agg. Benchmark #####
-    "overall": 86.0,
+    "overall": 65.92,
     ##### Code #####
     "code": 88.4,
 }
-# LLAMA3_2_11B_V_MODEL_CARD = {
-#     ##### Cost in USD #####
-#     "usd_per_input_token": 0.18 / 1E6,
-#     "usd_per_output_token": 0.18 / 1E6,
-#     ##### Time #####
-#     "seconds_per_output_token": 0.0061,
-#     ##### Agg. Benchmark #####
-#     "overall": 71.0,
-# }
 LLAMA3_2_90B_V_MODEL_CARD = {
     ##### Cost in USD #####
     "usd_per_input_token": 1.2 / 1e6,
@@ -258,7 +295,7 @@ LLAMA3_2_90B_V_MODEL_CARD = {
     ##### Time #####
     "seconds_per_output_token": 0.0222,
     ##### Agg. Benchmark #####
-    "overall": 84.0,
+    "overall": 65.00, # set to be slightly higher than gpt-4o-mini
 }
 MIXTRAL_8X_7B_MODEL_CARD = {
     ##### Cost in USD #####
@@ -267,7 +304,7 @@ MIXTRAL_8X_7B_MODEL_CARD = {
     ##### Time #####
     "seconds_per_output_token": 0.0112,
     ##### Agg. Benchmark #####
-    "overall": 63.0,
+    "overall": 43.27,
     ##### Code #####
     "code": 40.0,
 }
@@ -278,51 +315,56 @@ DEEPSEEK_V3_MODEL_CARD = {
     ##### Time #####
     "seconds_per_output_token": 0.0769,
     ##### Agg. Benchmark #####
-    "overall": 87.0,
+    "overall": 75.87,
     ##### Code #####
     "code": 92.0,
 }
-GPT_4o_MODEL_CARD = {
+DEEPSEEK_R1_DISTILL_QWEN_1_5B_MODEL_CARD = {
     ##### Cost in USD #####
-    "usd_per_input_token": 2.5 / 1e6,
-    "usd_per_output_token": 10.0 / 1e6,
+    "usd_per_input_token": 0.18 / 1E6,
+    "usd_per_output_token": 0.18 / 1E6,
     ##### Time #####
-    "seconds_per_output_token": 0.0079,
+    "seconds_per_output_token": 0.0026,
     ##### Agg. Benchmark #####
-    "overall": 89.0,
+    "overall": 39.90, # https://www.reddit.com/r/LocalLLaMA/comments/1iserf9/deepseek_r1_distilled_models_mmlu_pro_benchmarks/
     ##### Code #####
-    "code": 90.0,
+    "code": 0.0,
 }
-GPT_4o_V_MODEL_CARD = {
-    # NOTE: it is unclear if the same ($ / token) costs can be applied, or if we have to calculate this ourselves
+GPT_4o_MODEL_CARD = {
+    # NOTE: it is unclear if the same ($ / token) costs can be applied for vision, or if we have to calculate this ourselves
     ##### Cost in USD #####
     "usd_per_input_token": 2.5 / 1e6,
     "usd_per_output_token": 10.0 / 1e6,
     ##### Time #####
     "seconds_per_output_token": 0.0079,
     ##### Agg. Benchmark #####
-    "overall": 89.0,
+    "overall": 74.68,
+    ##### Code #####
+    "code": 90.0,
 }
 GPT_4o_MINI_MODEL_CARD = {
+    # NOTE: it is unclear if the same ($ / token) costs can be applied for vision, or if we have to calculate this ourselves
     ##### Cost in USD #####
     "usd_per_input_token": 0.15 / 1e6,
     "usd_per_output_token": 0.6 / 1e6,
     ##### Time #####
     "seconds_per_output_token": 0.0098,
     ##### Agg. Benchmark #####
-    "overall": 82.0,
+    "overall": 63.09,
     ##### Code #####
     "code": 86.0,
 }
-GPT_4o_MINI_V_MODEL_CARD = {
-    # NOTE: it is unclear if the same ($ / token) costs can be applied, or if we have to calculate this ourselves
+o1_MODEL_CARD = {  # noqa: N816
+    # NOTE: it is unclear if the same ($ / token) costs can be applied for vision, or if we have to calculate this ourselves
     ##### Cost in USD #####
-    "usd_per_input_token": 0.15 / 1e6,
-    "usd_per_output_token": 0.6 / 1e6,
+    "usd_per_input_token": 15 / 1e6,
+    "usd_per_output_token": 60 / 1e6,
     ##### Time #####
-    "seconds_per_output_token": 0.0098,
+    "seconds_per_output_token": 0.0110,
     ##### Agg. Benchmark #####
-    "overall": 82.0,
+    "overall": 89.30,
+    ##### Code #####
+    "code": 92.3, # NOTE: just copying MMLU score for now
 }
 TEXT_EMBEDDING_3_SMALL_MODEL_CARD = {
     ##### Cost in USD #####
@@ -331,7 +373,7 @@ TEXT_EMBEDDING_3_SMALL_MODEL_CARD = {
     ##### Time #####
     "seconds_per_output_token": 0.0098,  # NOTE: just copying GPT_4o_MINI_MODEL_CARD for now
     ##### Agg. Benchmark #####
-    "overall": 82.0,  # NOTE: just copying GPT_4o_MINI_MODEL_CARD for now
+    "overall": 63.09,  # NOTE: just copying GPT_4o_MINI_MODEL_CARD for now
 }
 CLIP_VIT_B_32_MODEL_CARD = {
     ##### Cost in USD #####
@@ -345,22 +387,18 @@ CLIP_VIT_B_32_MODEL_CARD = {
 
 
 MODEL_CARDS = {
-    Model.LLAMA3.value: LLAMA3_3_70B_INSTRUCT_MODEL_CARD,
-    Model.LLAMA3_V.value: LLAMA3_2_90B_V_MODEL_CARD,
-    Model.DEEPSEEK.value: DEEPSEEK_V3_MODEL_CARD,
+    Model.LLAMA3_2_3B.value: LLAMA3_2_3B_INSTRUCT_MODEL_CARD,
+    Model.LLAMA3_1_8B.value: LLAMA3_1_8B_INSTRUCT_MODEL_CARD,
+    Model.LLAMA3_3_70B.value: LLAMA3_3_70B_INSTRUCT_MODEL_CARD,
+    Model.LLAMA3_2_90B_V.value: LLAMA3_2_90B_V_MODEL_CARD,
+    Model.DEEPSEEK_V3.value: DEEPSEEK_V3_MODEL_CARD,
+    Model.DEEPSEEK_R1_DISTILL_QWEN_1_5B.value: DEEPSEEK_R1_DISTILL_QWEN_1_5B_MODEL_CARD,
     Model.MIXTRAL.value: MIXTRAL_8X_7B_MODEL_CARD,
     Model.GPT_4o.value: GPT_4o_MODEL_CARD,
-    Model.GPT_4o_V.value: GPT_4o_V_MODEL_CARD,
     Model.GPT_4o_MINI.value: GPT_4o_MINI_MODEL_CARD,
-    Model.GPT_4o_MINI_V.value: GPT_4o_MINI_V_MODEL_CARD,
+    # Model.o1.value: o1_MODEL_CARD,
     Model.TEXT_EMBEDDING_3_SMALL.value: TEXT_EMBEDDING_3_SMALL_MODEL_CARD,
     Model.CLIP_VIT_B_32.value: CLIP_VIT_B_32_MODEL_CARD,
-    ###
-    # Model.GPT_3_5.value: GPT_3_5_MODEL_CARD,
-    # Model.GPT_4.value: GPT_4_MODEL_CARD,
-    # Model.GPT_4V.value: GPT_4V_MODEL_CARD,
-    # Model.GEMINI_1.value: GEMINI_1_MODEL_CARD,
-    # Model.GEMINI_1V.value: GEMINI_1V_MODEL_CARD,
 }
 
 

palimpzest/core/data/dataclasses.py

@@ -5,6 +5,8 @@ from abc import abstractmethod
 from dataclasses import dataclass, field, fields
 from typing import Any
 
+import numpy as np
+
 
 @dataclass
 class GenerationStats:
@@ -148,7 +150,7 @@ class RecordOpStats:
     record_state: dict[str, Any]
 
     # operation id; an identifier for this operation's physical op id
-    op_id: str
+    full_op_id: str
 
     # logical operation id; the logical op id for this physical op
     logical_op_id: str
@@ -164,7 +166,7 @@ class RecordOpStats:
 
     ##### NOT-OPTIONAL, BUT FILLED BY EXECUTION CLASS AFTER CONSTRUCTOR CALL #####
     # the ID of the physical operation which produced the input record for this record at this operation
-    source_op_id: str | None = None
+    source_full_op_id: str | None = None
 
     # the ID of the physical plan which produced this record at this operation
     plan_id: str = ""
@@ -240,8 +242,8 @@ class OperatorStats:
     Dataclass for storing statistics captured within a given operator.
     """
 
-    # the ID of the physical operation in which these stats were collected
-    op_id: str
+    # the full ID of the physical operation in which these stats were collected
+    full_op_id: str
 
     # the name of the physical operation in which these stats were collected
     op_name: str
@@ -255,8 +257,8 @@ class OperatorStats:
     # a list of RecordOpStats processed by the operation
     record_op_stats_lst: list[RecordOpStats] = field(default_factory=list)
 
-    # the ID of the physical operator which precedes this one
-    source_op_id: str | None = None
+    # the full ID of the physical operator which precedes this one
+    source_full_op_id: str | None = None
 
     # the ID of the physical plan which this operator is part of
     plan_id: str = ""
@@ -273,7 +275,7 @@ class OperatorStats:
 
         NOTE: in case (1.) we assume the execution layer guarantees that `stats` is
               generated by the same operator in the same plan. Thus, we assume the
-              op_ids, op_name, source_op_id, etc. do not need to be updated.
+              full_op_ids, op_name, source_op_id, etc. do not need to be updated.
         """
         if isinstance(stats, OperatorStats):
             self.total_op_time += stats.total_op_time
@@ -281,7 +283,7 @@ class OperatorStats:
             self.record_op_stats_lst.extend(stats.record_op_stats_lst)
 
         elif isinstance(stats, RecordOpStats):
-            stats.source_op_id = self.source_op_id
+            stats.source_full_op_id = self.source_full_op_id
             stats.plan_id = self.plan_id
             self.record_op_stats_lst.append(stats)
             self.total_op_time += stats.time_per_record
@@ -294,7 +296,7 @@ class OperatorStats:
 
     def to_json(self):
         return {
-            "op_id": self.op_id,
+            "full_op_id": self.full_op_id,
             "op_name": self.op_name,
             "total_op_time": self.total_op_time,
             "total_op_cost": self.total_op_cost,
@@ -327,8 +329,8 @@ class BasePlanStats:
     plan_str: str | None = None
 
     # dictionary whose values are OperatorStats objects;
-    # PlanStats maps {physical_op_id -> OperatorStats}
-    # SentinelPlanStats maps {logical_op_id -> {physical_op_id -> OperatorStats}}
+    # PlanStats maps {full_op_id -> OperatorStats}
+    # SentinelPlanStats maps {logical_op_id -> {full_op_id -> OperatorStats}}
     operator_stats: dict = field(default_factory=dict)
 
     # total runtime for the plan measured from the start to the end of PhysicalPlan.execute()
@@ -406,11 +408,11 @@ class PlanStats(BasePlanStats):
         """
         operator_stats = {}
         for op_idx, op in enumerate(plan.operators):
-            op_id = op.get_op_id()
-            operator_stats[op_id] = OperatorStats(
-                op_id=op_id,
+            full_op_id = op.get_full_op_id()
+            operator_stats[full_op_id] = OperatorStats(
+                full_op_id=full_op_id,
                 op_name=op.op_name(),
-                source_op_id=None if op_idx == 0 else plan.operators[op_idx - 1].get_op_id(),
+                source_full_op_id=None if op_idx == 0 else plan.operators[op_idx - 1].get_full_op_id(),
                 plan_id=plan.plan_id,
                 op_details={k: str(v) for k, v in op.get_id_params().items()},
             )
@@ -432,11 +434,11 @@ class PlanStats(BasePlanStats):
 
         # update operator stats
         for record_op_stats in record_op_stats_lst:
-            op_id = record_op_stats.op_id
-            if op_id in self.operator_stats:
-                self.operator_stats[op_id] += record_op_stats
+            full_op_id = record_op_stats.full_op_id
+            if full_op_id in self.operator_stats:
+                self.operator_stats[full_op_id] += record_op_stats
             else:
-                raise ValueError(f"RecordOpStats with physical_op_id {op_id} not found in PlanStats")
+                raise ValueError(f"RecordOpStats with full_op_id {full_op_id} not found in PlanStats")
 
     def __iadd__(self, plan_stats: PlanStats) -> None:
         """
@@ -448,11 +450,11 @@ class PlanStats(BasePlanStats):
         """
         self.total_plan_time += plan_stats.total_plan_time
         self.total_plan_cost += plan_stats.total_plan_cost
-        for op_id, op_stats in plan_stats.operator_stats.items():
-            if op_id in self.operator_stats:
-                self.operator_stats[op_id] += op_stats
+        for full_op_id, op_stats in plan_stats.operator_stats.items():
+            if full_op_id in self.operator_stats:
+                self.operator_stats[full_op_id] += op_stats
             else:
-                self.operator_stats[op_id] = op_stats
+                self.operator_stats[full_op_id] = op_stats
 
     def __str__(self) -> str:
         stats = f"total_plan_time={self.total_plan_time} \n"
@@ -465,7 +467,7 @@ class PlanStats(BasePlanStats):
         return {
             "plan_id": self.plan_id,
             "plan_str": self.plan_str,
-            "operator_stats": {op_id: op_stats.to_json() for op_id, op_stats in self.operator_stats.items()},
+            "operator_stats": {full_op_id: op_stats.to_json() for full_op_id, op_stats in self.operator_stats.items()},
             "total_plan_time": self.total_plan_time,
             "total_plan_cost": self.total_plan_cost,
         }
@@ -485,11 +487,11 @@ class SentinelPlanStats(BasePlanStats):
         for op_set_idx, (logical_op_id, op_set) in enumerate(plan):
             operator_stats[logical_op_id] = {}
             for physical_op in op_set:
-                op_id = physical_op.get_op_id()
-                operator_stats[logical_op_id][op_id] = OperatorStats(
-                    op_id=op_id,
+                full_op_id = physical_op.get_full_op_id()
+                operator_stats[logical_op_id][full_op_id] = OperatorStats(
+                    full_op_id=full_op_id,
                     op_name=physical_op.op_name(),
-                    source_op_id=None if op_set_idx == 0 else plan.logical_op_ids[op_set_idx - 1],
+                    source_full_op_id=None if op_set_idx == 0 else plan.logical_op_ids[op_set_idx - 1],  # NOTE: this may be a reason to keep `source_op_id` instead of `source_full_op_id`
                     plan_id=plan.plan_id,
                     op_details={k: str(v) for k, v in physical_op.get_id_params().items()},
                 )
@@ -512,12 +514,12 @@ class SentinelPlanStats(BasePlanStats):
         # update operator stats
         for record_op_stats in record_op_stats_lst:
             logical_op_id = record_op_stats.logical_op_id
-            physical_op_id = record_op_stats.op_id
+            full_op_id = record_op_stats.full_op_id
             if logical_op_id in self.operator_stats:
-                if physical_op_id in self.operator_stats[logical_op_id]:
-                    self.operator_stats[logical_op_id][physical_op_id] += record_op_stats
+                if full_op_id in self.operator_stats[logical_op_id]:
+                    self.operator_stats[logical_op_id][full_op_id] += record_op_stats
                 else:
-                    raise ValueError(f"RecordOpStats with physical_op_id {physical_op_id} not found in SentinelPlanStats")
+                    raise ValueError(f"RecordOpStats with full_op_id {full_op_id} not found in SentinelPlanStats")
             else:
                 raise ValueError(f"RecordOpStats with logical_op_id {logical_op_id} not found in SentinelPlanStats")
 
@@ -532,12 +534,12 @@ class SentinelPlanStats(BasePlanStats):
         self.total_plan_time += plan_stats.total_plan_time
         self.total_plan_cost += plan_stats.total_plan_cost
         for logical_op_id, physical_op_stats in plan_stats.operator_stats.items():
-            for physical_op_id, op_stats in physical_op_stats.items():
+            for full_op_id, op_stats in physical_op_stats.items():
                 if logical_op_id in self.operator_stats:
-                    if physical_op_id in self.operator_stats[logical_op_id]:
-                        self.operator_stats[logical_op_id][physical_op_id] += op_stats
+                    if full_op_id in self.operator_stats[logical_op_id]:
+                        self.operator_stats[logical_op_id][full_op_id] += op_stats
                     else:
-                        self.operator_stats[logical_op_id][physical_op_id] = op_stats
+                        self.operator_stats[logical_op_id][full_op_id] = op_stats
                 else:
                     self.operator_stats[logical_op_id] = physical_op_stats
 
@@ -557,7 +559,7 @@ class SentinelPlanStats(BasePlanStats):
             "plan_id": self.plan_id,
             "plan_str": self.plan_str,
             "operator_stats": {
-                logical_op_id: {physical_op_id: op_stats.to_json() for physical_op_id, op_stats in physical_op_stats.items()}
+                logical_op_id: {full_op_id: op_stats.to_json() for full_op_id, op_stats in physical_op_stats.items()}
                 for logical_op_id, physical_op_stats in self.operator_stats.items()
             },
             "total_plan_time": self.total_plan_time,
@@ -684,8 +686,21 @@ class ExecutionStats:
             else:
                 raise TypeError(f"Cannot add {type(plan_stats)} to ExecutionStats")
 
+    def clean_json(self, stats: dict):
+        """
+        Convert np.int64 and np.float64 to int and float for all values in stats.
+        """
+        for key, value in stats.items():
+            if isinstance(value, dict):
+                stats[key] = self.clean_json(value)
+            elif isinstance(value, np.int64):
+                stats[key] = int(value)
+            elif isinstance(value, np.float64):
+                stats[key] = float(value)
+        return stats
+
     def to_json(self):
-        return {
+        stats = {
             "execution_id": self.execution_id,
             "sentinel_plan_stats": {
                 plan_id: plan_stats.to_json() for plan_id, plan_stats in self.sentinel_plan_stats.items()
@@ -700,6 +715,8 @@ class ExecutionStats:
             "sentinel_plan_strs": self.sentinel_plan_strs,
             "plan_strs": self.plan_strs,
         }
+        stats = self.clean_json(stats)
+        return stats
 
 
 @dataclass

palimpzest/core/elements/index.py

@@ -3,30 +3,27 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 
 from chromadb.api.models.Collection import Collection
-from ragatouille.RAGPretrainedModel import RAGPretrainedModel
 
 
-def index_factory(index: Collection | RAGPretrainedModel) -> PZIndex:
+def index_factory(index: Collection) -> PZIndex:
     """
     Factory function to create a PZ index based on the type of the provided index.
 
     Args:
-        index (Collection | RAGPretrainedModel): The index provided by the user.
+        index (Collection): The index provided by the user.
 
     Returns:
         PZIndex: The PZ wrapped Index.
     """
     if isinstance(index, Collection):
         return ChromaIndex(index)
-    elif isinstance(index, RAGPretrainedModel):
-        return RagatouilleIndex(index)
     else:
-        raise TypeError(f"Unsupported index type: {type(index)}\nindex must be a `chromadb.api.models.Collection.Collection` or `ragatouille.RAGPretrainedModel.RAGPretrainedModel`")
+        raise TypeError(f"Unsupported index type: {type(index)}\nindex must be a `chromadb.api.models.Collection.Collection`")
 
 
 class BaseIndex(ABC):
 
-    def __init__(self, index: Collection | RAGPretrainedModel):
+    def __init__(self, index: Collection):
         self.index = index
 
     def __str__(self):
@@ -59,12 +56,5 @@ class ChromaIndex(BaseIndex):
         super().__init__(index)
 
 
-
-class RagatouilleIndex(BaseIndex):
-    def __init__(self, index: RAGPretrainedModel):
-        assert isinstance(index, RAGPretrainedModel), "RagatouilleIndex input must be a `ragatouille.RAGPretrainedModel.RAGPretrainedModel`"
-        super().__init__(index)
-
-
 # define type for PZIndex
-PZIndex = ChromaIndex | RagatouilleIndex
+PZIndex = ChromaIndex

palimpzest/core/elements/records.py

@@ -35,7 +35,7 @@ class DataRecord:
         self.field_values: dict[str, Any] = {}
 
         # the index in the DataReader from which this DataRecord is derived
-        self.source_idx = source_idx
+        self.source_idx = int(source_idx)
 
         # the id of the parent record(s) from which this DataRecord is derived
         self.parent_id = parent_id

palimpzest/prompts/prompt_factory.py

@@ -155,7 +155,7 @@ class PromptFactory:
         # TODO: this does not work for image prompts
         # TODO: this ignores the size of the `orignal_messages` in critique and refine prompts
         # cut down on context based on window length
-        if self.model in [Model.LLAMA3, Model.MIXTRAL]:
+        if self.model.is_llama_model() or self.model.is_mixtral_model():
             total_context_len = len(json.dumps(context, indent=2))
 
             # sort fields by length and progressively strip from the longest field until it is short enough;