feat: save factor score in the trace (#1350)

Signed-off-by: Louis Mandel <[email protected]>

Author: mandel

pdl-live-react/src/pdl_ast.d.ts

@@ -447,7 +447,15 @@ export type Contribute19 = ContributeElement[]
 export type ExpectationsType19 = ExpectationType[]
 export type PdlIsLeaf19 = true
 export type Kind19 = "factor"
+/**
+ * Score to condition the model.
+ *
+ */
 export type Factor = LocalizedExpression | number | string
+/**
+ * Allow to raise the Resampling exception during the execution of the block.
+ *
+ */
 export type Resample = boolean
 /**
  * Indicate if the block contributes to the result and background context.
@@ -3368,6 +3376,7 @@ export interface FactorBlock {
   kind?: Kind19
   factor: Factor
   resample?: Resample
+  pdl__score?: number | null
 }
 /**
  * Set of definitions executed before the execution of the block.

src/pdl/pdl-schema.json

@@ -1425,12 +1425,18 @@
               "type": "string"
             }
           ],
+          "description": "Score to condition the model.\n    ",
           "title": "Factor"
         },
         "resample": {
           "default": true,
+          "description": "Allow to raise the Resampling exception during the execution of the block.\n    ",
           "title": "Resample",
           "type": "boolean"
+        },
+        "pdl__score": {
+          "$ref": "#/$defs/OptionalFloat",
+          "default": null
         }
       },
       "required": [
@@ -4029,6 +4035,16 @@
         }
       ]
     },
+    "OptionalFloat": {
+      "anyOf": [
+        {
+          "type": "number"
+        },
+        {
+          "type": "null"
+        }
+      ]
+    },
     "OptionalInt": {
       "anyOf": [
         {

src/pdl/pdl_ast.py

@@ -41,6 +41,8 @@ def _ensure_lower(value):
 """Optional string."""
 OptionalInt = TypeAliasType("OptionalInt", Optional[int])
 """Optional integer."""
+OptionalFloat = TypeAliasType("OptionalFloat", Optional[float])
+"""Optional floating point number."""
 OptionalBool = TypeAliasType("OptionalBool", Optional[bool])
 """Optional Boolean."""
 OptionalBoolOrStr = TypeAliasType("OptionalBoolOrStr", Optional[Union[bool, str]])
@@ -1093,7 +1095,13 @@ class FactorBlock(LeafBlock):
 
     kind: Literal[BlockKind.FACTOR] = BlockKind.FACTOR
     factor: ExpressionType[float]
+    """Score to condition the model.
+    """
     resample: bool = True
+    """Allow to raise the Resampling exception during the execution of the block.
+    """
+    # Field for internal use
+    pdl__score: OptionalFloat = None
 
 
 class AggregatorConfig(BaseModel):

src/pdl/pdl_dumper.py

@@ -229,6 +229,8 @@ def block_to_dict(  # noqa: C901
             d["factor"] = expr_to_dict(block.factor, json_compatible)
             if not block.resample:
                 d["resample"] = block.resample
+            if block.pdl__score is not None:
+                d["pdl__score"] = block.pdl__score
         case AggregatorBlock():
             d["aggregator"] = aggregator_to_dict(block.aggregator)
         case IfBlock():

src/pdl/pdl_interpreter.py

@@ -1138,6 +1138,7 @@ def loop_body(iidx, items):
                 weight, trace = process_expr_of(
                     block, "factor", scope, append(loc, "factor")
                 )
+            trace.pdl__score = weight
             state.score.ref += weight
             result = PdlConst("")
             background = DependentContext([])