datago 2025.3.11__tar.gz → 2025.4.1__tar.gz

{datago-2025.3.11 → datago-2025.4.1}/Cargo.lock

@@ -387,7 +387,7 @@ dependencies = [
 
 [[package]]
 name = "datago"
-version = "2025.3.11"
+version = "2025.4.1"
 dependencies = [
  "clap",
  "image",

{datago-2025.3.11 → datago-2025.4.1}/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "datago"
 edition = "2021"
-version = "2025.3.11"
+version = "2025.4.1"
 
 [lib]
 # exposed by pyo3
@@ -14,7 +14,7 @@ name = "datago"
 path = "src/main.rs"
 
 [dependencies]
-image = "0.25.5"
+image = { version = "0.25.5",  features = ["png"]}
 reqwest = { version = "0.12.12", features = ["blocking"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"

{datago-2025.3.11 → datago-2025.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datago
-Version: 2025.3.11
+Version: 2025.4.1
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
@@ -102,6 +102,23 @@ See helper functions provided in `raw_types.py`, should be self explanatory. Che
 
 Just install the rust toolchain via rustup
 
+## [Apple Silicon MacOS only]
+
+If you are using an Apple Silicon Mac OS machine, create a `.cargo/config` file and paste the following:
+```
+[target.x86_64-apple-darwin]
+rustflags = [
+  "-C", "link-arg=-undefined",
+  "-C", "link-arg=dynamic_lookup",
+]
+
+[target.aarch64-apple-darwin]
+rustflags = [
+  "-C", "link-arg=-undefined",
+  "-C", "link-arg=dynamic_lookup",
+]
+```
+
 ## Build a benchmark CLI
 `cargo run --release --  -h` to get all the information, should be fairly straightforward
 

{datago-2025.3.11 → datago-2025.4.1}/README.md

@@ -86,6 +86,23 @@ See helper functions provided in `raw_types.py`, should be self explanatory. Che
 
 Just install the rust toolchain via rustup
 
+## [Apple Silicon MacOS only]
+
+If you are using an Apple Silicon Mac OS machine, create a `.cargo/config` file and paste the following:
+```
+[target.x86_64-apple-darwin]
+rustflags = [
+  "-C", "link-arg=-undefined",
+  "-C", "link-arg=dynamic_lookup",
+]
+
+[target.aarch64-apple-darwin]
+rustflags = [
+  "-C", "link-arg=-undefined",
+  "-C", "link-arg=dynamic_lookup",
+]
+```
+
 ## Build a benchmark CLI
 `cargo run --release --  -h` to get all the information, should be fairly straightforward
 

{datago-2025.3.11 → datago-2025.4.1}/src/generator_http.rs

@@ -24,9 +24,18 @@ pub struct SourceDBConfig {
     #[serde(default)]
     pub tags: String,
 
+    #[serde(default)]
+    pub tags_all: String,
+
     #[serde(default)]
     pub tags_ne: String,
 
+    #[serde(default)]
+    pub tags_ne_all: String,
+
+    #[serde(default)]
+    pub tags_empty: String,
+
     #[serde(default)]
     pub has_attributes: String,
 
@@ -60,6 +69,9 @@ pub struct SourceDBConfig {
     #[serde(default)]
     pub duplicate_state: i32,
 
+    #[serde(default)]
+    pub attributes: String,
+
     #[serde(default)]
     pub random_sampling: bool,
 }
@@ -73,7 +85,10 @@ struct DbRequest {
     pub page_size: String,
 
     pub tags: String,
+    pub tags_all: String,
     pub tags_ne: String,
+    pub tags_ne_all: String,
+    pub tags_empty: String,
 
     pub has_attributes: String,
     pub lacks_attributes: String,
@@ -92,6 +107,7 @@ struct DbRequest {
     pub max_pixel_count: String,
 
     pub duplicate_state: String,
+    pub attributes: String,
     pub random_sampling: bool,
 
     pub partitions_count: String,
@@ -131,7 +147,10 @@ impl DbRequest {
             maybe_add_field("page_size", &self.page_size);
 
             maybe_add_field("tags", &self.tags);
+            maybe_add_field("tags__all", &self.tags_all);
             maybe_add_field("tags__ne", &self.tags_ne);
+            maybe_add_field("tags__ne_all", &self.tags_ne_all);
+            maybe_add_field("tags__empty", &self.tags_empty);
             maybe_add_field("has_attributes", &self.has_attributes);
             maybe_add_field("lacks_attributes", &self.lacks_attributes);
             maybe_add_field("has_masks", &self.has_masks);
@@ -144,6 +163,7 @@ impl DbRequest {
             maybe_add_field("pixel_count__gte", &self.min_pixel_count);
             maybe_add_field("pixel_count__lte", &self.max_pixel_count);
             maybe_add_field("duplicate_state", &self.duplicate_state);
+            maybe_add_field("attributes", &self.attributes);
             maybe_add_field("partitions_count", &self.partitions_count);
             maybe_add_field("partition", &self.partition);
         }
@@ -180,7 +200,53 @@ fn build_request(source_config: SourceDBConfig, rank: usize, world_size: usize)
 
     if !source_config.tags.is_empty() {
         fields.push_str(",tags");
-        println!("Including some tags: {}", source_config.tags);
+        println!(
+            "Including some tags, must have any of: {}",
+            source_config.tags
+        );
+    }
+
+    if !source_config.tags_all.is_empty() {
+        fields.push_str(",tags");
+        println!(
+            "Including tags, must have all of: {}",
+            source_config.tags_all
+        );
+    }
+
+    if !source_config.tags_ne.is_empty() {
+        fields.push_str(",tags");
+        println!(
+            "Including tags, must not have any of: {}",
+            source_config.tags_ne
+        );
+    }
+
+    if !source_config.tags_empty.is_empty() {
+        fields.push_str(",tags");
+        println!(
+            "Using filter: Tags must{} be empty",
+            if source_config.tags_empty == "true" {
+                " not"
+            } else {
+                ""
+            }
+        );
+        if !source_config.tags_all.is_empty()
+            || !source_config.tags.is_empty()
+            || !source_config.tags_ne.is_empty()
+            || !source_config.tags_ne_all.is_empty()
+        {
+            println!("WARNING: you've set `tags_empty` in addition to `tags`, `tags_all`, `tags_ne` or `tags_ne_all`. The combination might be incompatible or redundant.");
+        }
+    }
+
+    if !source_config.tags_ne_all.is_empty() {
+        fields.push_str(",tags");
+        println!(
+            "Including tags, must not have all of: {}",
+            source_config.tags_ne_all
+        );
     }
 
     if source_config.require_embeddings {
@@ -213,7 +279,10 @@ fn build_request(source_config: SourceDBConfig, rank: usize, world_size: usize)
         sources_ne: source_config.sources_ne,
         page_size: source_config.page_size.to_string(),
         tags: source_config.tags,
+        tags_all: source_config.tags_all,
         tags_ne: source_config.tags_ne,
+        tags_ne_all: source_config.tags_ne_all,
+        tags_empty: source_config.tags_empty,
         has_attributes: source_config.has_attributes,
         lacks_attributes: source_config.lacks_attributes,
         has_masks: source_config.has_masks,
@@ -226,6 +295,7 @@ fn build_request(source_config: SourceDBConfig, rank: usize, world_size: usize)
         min_pixel_count: maybe_add_int(source_config.min_pixel_count),
         max_pixel_count: maybe_add_int(source_config.max_pixel_count),
         duplicate_state: maybe_add_int(source_config.duplicate_state),
+        attributes: source_config.attributes,
         random_sampling: source_config.random_sampling,
         partition: if world_size > 1 {
             format!("{}", rank)

{datago-2025.3.11 → datago-2025.4.1}/src/image_processing.rs

@@ -1,8 +1,8 @@
 use crate::structs::ImagePayload;
+use image::ImageEncoder;
 use serde::Deserialize;
 use serde::Serialize;
 use std::io::Cursor;
-
 // --- Sample data structures - these will be exposed to the Python world ---------------------------------------------------------------------------------------------------------------------------------------------------------------
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -176,15 +176,24 @@ pub async fn image_to_payload(
     // Encode the image if needed
     let mut image_bytes: Vec<u8> = Vec::new();
     if encode_images {
-        if image
-            .write_to(&mut Cursor::new(&mut image_bytes), image::ImageFormat::Png)
-            .is_err()
-        {
-            return Err(image::ImageError::IoError(std::io::Error::new(
+        // Use the encoder directly with the raw bytes
+        image::codecs::png::PngEncoder::new_with_quality(
+            &mut Cursor::new(&mut image_bytes),
+            image::codecs::png::CompressionType::Fast,
+            image::codecs::png::FilterType::Adaptive,
+        )
+        .write_image(
+            image.as_bytes(),
+            image.width(),
+            image.height(),
+            image.color().into(),
+        )
+        .map_err(|e| {
+            image::ImageError::IoError(std::io::Error::new(
                 std::io::ErrorKind::Other,
-                "Failed to encode image",
-            )));
-        }
+                e.to_string(),
+            ))
+        })?;
 
         channels = -1; // Signal the fact that the image is encoded
     } else {

{datago-2025.3.11 → datago-2025.4.1}/tests/client_test.rs

@@ -19,6 +19,9 @@ fn get_test_config() -> serde_json::Value {
             "require_embeddings": false,
             "tags": "",
             "tags_ne": "",
+            "tags_all": "",
+            "tags_ne_all": "",
+            "tags_empty": "",
             "has_attributes": "",
             "lacks_attributes": "",
             "has_masks": "",
@@ -30,6 +33,7 @@ fn get_test_config() -> serde_json::Value {
             "min_pixel_count": -1,
             "max_pixel_count": -1,
             "duplicate_state": -1,
+            "attributes": "",
             "random_sampling": false,
             "page_size": 10,
         },
@@ -253,6 +257,116 @@ fn test_tags() {
     client.stop();
 }
 
+#[test]
+fn test_tags_all() {
+    let mut config = get_test_config();
+    let tags = "v4_trainset_hq,photo";
+    config["source_config"]["tags_all"] = tags.into();
+    let mut client = DatagoClient::new(config.to_string());
+
+    let sample = client.get_sample();
+    assert!(sample.is_some());
+
+    let sample = sample.unwrap();
+    assert!(!sample.id.is_empty());
+    // Check that sample.tags contains all the tags in the tags string
+    for tag in tags.split(',') {
+        assert!(sample.tags.contains(&tag.to_string()));
+    }
+    client.stop();
+}
+
+#[test]
+fn test_tags_ne() {
+    let mut config = get_test_config();
+    let tags = "v4_trainset_hq,photo";
+    config["source_config"]["tags_ne"] = tags.into();
+    let mut client = DatagoClient::new(config.to_string());
+
+    let sample = client.get_sample();
+    assert!(sample.is_some());
+
+    let sample = sample.unwrap();
+    assert!(!sample.id.is_empty());
+    // Check that sample.tags does not contain any of the tags in the tags string
+    println!("{:?}", sample.tags);
+    for tag in tags.split(',') {
+        assert!(!sample.tags.contains(&tag.to_string()));
+    }
+    client.stop();
+}
+
+#[test]
+fn test_tags_empty() {
+    let mut config = get_test_config();
+    config["source_config"]["tags_empty"] = "true".into();
+    let mut client = DatagoClient::new(config.to_string());
+
+    let sample = client.get_sample();
+    assert!(sample.is_some());
+
+    let sample = sample.unwrap();
+    assert!(sample.tags.is_empty());
+    client.stop();
+}
+
+#[test]
+fn test_tags_ne_all() {
+    let mut config = get_test_config();
+    let tag1 = "photo";
+    let tag2 = "graphic";
+    config["source_config"]["tags_ne_all"] = format!("{},{}", tag1, tag2).into();
+    let mut client = DatagoClient::new(config.to_string());
+
+    let sample = client.get_sample();
+    assert!(sample.is_some());
+
+    let sample = sample.unwrap();
+    assert!(!sample.id.is_empty());
+    // Assert that the sample does not contain both tags at the same time
+    let has_first = sample.tags.contains(&tag1.to_string());
+    let has_second = sample.tags.contains(&tag2.to_string());
+    assert!(
+        !(has_first && has_second),
+        "Sample should not contain both tags at the same time"
+    );
+    client.stop();
+}
+
+#[test]
+fn test_attributes_filter() {
+    let mut config = get_test_config();
+    config["source_config"]["attributes"] = "aesthetic_score__gte:0.5".into();
+    let mut client = DatagoClient::new(config.to_string());
+
+    let sample = client.get_sample();
+    assert!(sample.is_some());
+
+    let sample = sample.unwrap();
+    assert!(!sample.id.is_empty());
+    assert!(sample.attributes.contains_key("aesthetic_score"));
+    assert!(sample.attributes["aesthetic_score"].as_f64().unwrap() >= 0.5);
+    client.stop();
+}
+
+#[test]
+fn test_pixel_count_filter() {
+    let mut config = get_test_config();
+    config["source_config"]["min_pixel_count"] = 1000000.into();
+    config["source_config"]["max_pixel_count"] = 2000000.into();
+    config["source_config"]["require_images"] = json!(true);
+    let mut client = DatagoClient::new(config.to_string());
+
+    let sample = client.get_sample();
+    assert!(sample.is_some());
+
+    let sample = sample.unwrap();
+    assert!(!sample.id.is_empty());
+    assert!(sample.image.width * sample.image.height >= 1000000);
+    assert!(sample.image.width * sample.image.height <= 2000000);
+    client.stop();
+}
+
 #[test]
 fn test_multiple_sources() {
     let limit = 10;