Polymorphic span replacement

Makefile

@@ -23,7 +23,7 @@ setup:
 	$(shell "${PROTOBUF_SETUP}")
 	$(shell "${OPENSSL_SETUP}")
 	which cargo || curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
-	which maturin || pip install maturin[patchelf]
+	which maturin || pip install maturin
 
 publish:
 	maturin publish

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dolma"
-version = "1.0.12"
+version = "1.0.13"
 description = "Data filters"
 license = { text = "Apache-2.0" }
 readme = "README.md"

python/dolma/cli/mixer.py

@@ -40,7 +40,7 @@ class SpanReplacementConfig:
         default=None,
         help="Maximum score for the span to be replaced. Either min_score or max_score must be specified.",
     )
-    replacement: str = field(default="", help="Replacement for the span")
+    replacement: str = field(default="", help="Replacement config for the span(s).")
     syntax: str = field(
         default="jsonpath",
         help="Syntax to use for filter expressions. Currently only JSONPath is supported. Defaults to JSONPath.",

src/shard.rs

@@ -399,9 +399,10 @@ impl Shard {
                                     );
                                 new_text.push_str(&replacement_text);
                             }
+
                             data["text"] = Value::String(new_text);
                         }
-                        // }
+
                         for f in self.discard_fields.iter().flatten() {
                             data.as_object_mut().unwrap().remove(f);
                         }
@@ -467,7 +468,7 @@ impl Shard {
 }
 
 pub mod shard_config {
-    use crate::filters::Selector;
+    use crate::filters::{JqSelector, Selector};
     use jsonpath_rust::JsonPathFinder;
     use serde::{Deserialize, Serialize};
     use serde_json::Value;
@@ -543,7 +544,34 @@ pub mod shard_config {
         selector: Selector,
         min_score: f64,
         max_score: f64,
-        replacement: String,
+        replacement: Replacement,
+    }
+
+    pub enum Replacement {
+        Selectors(JqSelector),
+        String(String),
+    }
+
+    impl Replacement {
+        pub fn new(string: &str) -> Result<Replacement, IoError> {
+            // Note: Users should escape leading $ in replacement strings
+            if string.starts_with("$") {
+                // Strip leading $ and create a selector
+                let selector = JqSelector::new(&string[1..])?;
+                Ok(Replacement::Selectors(selector))
+            } else {
+                Ok(Replacement::String(string.to_string()))
+            }
+        }
+
+        pub fn get(&self, json: &Value) -> Result<String, IoError> {
+            match self {
+                Replacement::Selectors(selector) => {
+                    Ok(serde_json::from_value(selector.select(json)?.to_owned()).unwrap())
+                }
+                Replacement::String(s) => Ok(s.clone()),
+            }
+        }
     }
 
     impl SpanReplacer {
@@ -553,7 +581,7 @@ pub mod shard_config {
                 selector: Selector::new(&config).unwrap(),
                 min_score: config.min_score.unwrap_or(f64::NEG_INFINITY),
                 max_score: config.max_score.unwrap_or(f64::INFINITY),
-                replacement: config.replacement.clone(),
+                replacement: Replacement::new(&config.replacement).unwrap(),
             }
         }
 
@@ -575,7 +603,7 @@ pub mod shard_config {
                                 let replacement = SpanReplacement {
                                     start: start as usize,
                                     end: end as usize,
-                                    replacement: self.replacement.clone(),
+                                    replacement: self.replacement.get(json).unwrap(),
                                 };
                                 Some(replacement)
                             } else {