Merge pull request #183 from allenai/attr-merge

Preserve existing attributes
allenai · May 26, 2023 · 72d4ff2 · 72d4ff2
2 parents 47c4ab9 + e0d8fc1
commit 72d4ff2
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/pretrain_data/mixer/src/shard.rs b/pretrain_data/mixer/src/shard.rs
@@ -221,7 +221,14 @@ impl Shard {
                     }
 
                     if !attrs.is_empty() {
-                        mutable_data["attributes"] = Value::Object(attrs);
+                        // Add to existing attributes if they exist, otherwise create them.
+                        if let Value::Object(ref mut existing_attrs) = mutable_data["attributes"] {
+                            for (k, v) in attrs.iter() {
+                                existing_attrs.insert(k.clone(), v.clone());
+                            }
+                        } else {
+                            mutable_data["attributes"] = Value::Object(attrs);
+                        }
                     }
 
                     let mut should_write = true;