Argilla datasets

Opt a protobuf message into Argilla codegen with (pgml.argilla).enable = true. py-gen-ml emits:

• Pydantic models with Field(description=...) from proto comments
• build_*_settings() for Argilla Settings (UI fields + questions)
• to_*_record / from_*_record mappers
• Optional alias helpers when (pgml.mapper_config).enable is set

Use message kinds (FEATURE_ROW, LABEL, FEEDBACK, …) to document the contract. Kind does not replace the Argilla opt-in.

Field vs Question vs Metadata

Every field on an Argilla-enabled message must set (pgml.argilla_field).slot. Missing slots fail at codegen.

Slot	Argilla artifact	Record mapping
FIELD	TextField / ChatField / ImageField / …	record.fields[name]
QUESTION	LabelQuestion / MultiLabelQuestion / RatingQuestion / TextQuestion / …	suggestions / responses
METADATA	(none in Settings)	record.metadata[name]

Typical pattern for classification HITL:

• FIELD — the text (or multimodal content) annotators read
• QUESTION — the label they assign (positive / negative, …)
• METADATA — ids, provenance (imdb vs synthetic), timestamps

string id = 1 [(pgml.argilla_field) = { slot: METADATA }];
string text = 2 [(pgml.argilla_field) = {
  slot: FIELD; field_type: "text"; required: true
}];
string sentiment = 3 [(pgml.argilla_field) = {
  slot: QUESTION; question_type: "label";
  labels: ["negative", "positive"]; required: true
}];

ArgillaField knobs

Field	Applies to	Meaning
slot	all	FIELD / QUESTION / METADATA (required)
name	all	Optional UI/API name (default: proto field name)
field_type	FIELD	text (default), chat, image, …
question_type	QUESTION	label, multi_label, rating, text, …
labels	label questions	Allowed class names
required	FIELD / QUESTION	Passed through to Argilla

Install the extra

pip install 'py-gen-ml[argilla]'

uv add 'py-gen-ml[argilla]'

The Argilla extra also pins datasets>=3 so imports work with modern pyarrow (versions that removed PyExtensionType, commonly pulled in with LanceDB).

py-gen-ml path/to/schema.proto --generators=base,patch,sweep,cli_args,argilla

Annotate

Instruction–response style demo:

syntax = "proto3";

package demo;

import "py_gen_ml/extensions.proto";

// A single instruction–response example for LLM synthesis and Argilla review.
message ReviewExample {
    option (pgml.kind) = FEATURE_ROW;
    option (pgml.pydantic_ai) = {
        enable: true;
    };
    option (pgml.argilla) = {
        enable: true;
        dataset_name: "review_examples";
    };
    option (pgml.mapper_config) = {
        enable: true;
    };

    // Stable identifier stored as Argilla metadata (not a primary UI field).
    string id = 1 [(pgml.argilla_field) = {
        slot: METADATA;
    }];

    // User-facing task prompt the model should answer. Clear comments become
    // JSON Schema descriptions used by PydanticAI NativeOutput.
    string instruction = 2 [(pgml.argilla_field) = {
        slot: FIELD;
        field_type: "text";
        required: true;
    }];

    // Complete answer or generation for the instruction.
    string generation = 3 [(pgml.argilla_field) = {
        slot: FIELD;
        field_type: "text";
        required: true;
    }];

    // Human quality judgment for the generation (Argilla LabelQuestion).
    string quality = 4 [(pgml.argilla_field) = {
        slot: QUESTION;
        question_type: "label";
        labels: ["bad", "good"];
        required: true;
    }];
}

Sentiment / IMDB flywheel contract:

syntax = "proto3";

package demo;

import "py_gen_ml/extensions.proto";

// Labeled movie-review example for sentiment classification.
// Field comments become JSON Schema descriptions for PydanticAI NativeOutput.
message SentimentExample {
    option (pgml.kind) = FEATURE_ROW;
    option (pgml.pydantic_ai) = {
        enable: true;
    };
    option (pgml.argilla) = {
        enable: true;
        dataset_name: "imdb_sentiment";
    };
    option (pgml.lancedb) = {
        enable: true;
        table_name: "sentiment_examples";
    };
    option (pgml.mapper_config) = {
        enable: true;
    };

    // Stable row id (Argilla metadata; not a primary UI field).
    string id = 1 [
        (pgml.argilla_field) = { slot: METADATA; },
        (pgml.lancedb_field) = { merge_key: true }
    ];

    // Provenance of the row: "imdb" for seed reviews, "synthetic" for LLM-generated ones.
    string source = 2 [(pgml.argilla_field) = {
        slot: METADATA;
    }];

    // Full movie-review text in the style of IMDB user reviews (may include spoilers,
    // informal tone, and mixed praise/criticism). Used as the classifier input.
    string text = 3 [(pgml.argilla_field) = {
        slot: FIELD;
        field_type: "text";
        required: true;
    }];

    // Binary sentiment of the review: "negative" or "positive".
    string sentiment = 4 [(pgml.argilla_field) = {
        slot: QUESTION;
        question_type: "label";
        labels: ["negative", "positive"];
        required: true;
    }];
}

// Hyperparameters for a sentiment flywheel training run.
message SentimentTrainConfig {
    option (pgml.kind) = RUN_CONFIG;
    option (pgml.cli) = {
        enable: true;
    };
    option (pgml.mlflow) = {
        enable: true;
        experiment_name: "imdb_sentiment";
        run_name_field: "run_name";
    };

    // Human-readable run name (logged as a tracker tag).
    string run_name = 1 [
        (pgml.default).string = "sentiment",
        (pgml.tracking_field) = { slot: TAG }
    ];
    // Number of synthetic examples requested per diversify round.
    int32 synthesize_count = 2 [(pgml.default).int32 = 4];
    // Extra synthesis rounds that feed prior outputs back as examples.
    int32 diversify_rounds = 3 [(pgml.default).int32 = 1];
    // OpenAI / Azure deployment name used for synthesis.
    string openai_model = 4 [(pgml.default).string = "gpt-4o"];
    // TF-IDF + logistic regression holdout fraction.
    float test_size = 5 [(pgml.default).float = 0.25];
}

// Holdout metrics from the sentiment classifier.
message SentimentMetrics {
    option (pgml.kind) = METRIC_SET;
    option (pgml.mlflow) = {
        enable: true;
    };

    float accuracy = 1;
    int32 n_train = 2;
    int32 n_test = 3;
    int32 n_labeled = 4;
}

// Inference request for the online sentiment classifier.
message SentimentPredictRequest {
    option (pgml.kind) = FEATURE_ROW;

    // Stable sample id (links prediction and feedback rows).
    string id = 1;
    // Review text to classify.
    string text = 2;
}

// Model output for one scored review (separate from the labeled training row).
message SentimentPrediction {
    option (pgml.kind) = PREDICTION;
    option (pgml.lancedb) = {
        enable: true;
        table_name: "sentiment_predictions";
    };

    // Sample id matching SentimentPredictRequest.id.
    string sample_id = 1 [(pgml.lancedb_field) = { merge_key: true }];
    // Review text that was scored.
    string text = 2;
    // Predicted sentiment: "negative" or "positive".
    string sentiment = 3;
    // Model confidence for the predicted class.
    float score = 4;
    // Deployed model / pipeline version tag.
    string model_version = 5;
}

// Human correction / re-label for a scored review.
message SentimentFeedback {
    option (pgml.kind) = FEEDBACK;
    option (pgml.argilla) = {
        enable: true;
        dataset_name: "imdb_sentiment_feedback";
    };
    option (pgml.lancedb) = {
        enable: true;
        table_name: "sentiment_feedback";
    };

    // Sample id matching SentimentPrediction.sample_id.
    string sample_id = 1 [
        (pgml.argilla_field) = { slot: METADATA; },
        (pgml.lancedb_field) = { merge_key: true }
    ];

    // Review text shown to the annotator.
    string text = 2 [(pgml.argilla_field) = {
        slot: FIELD;
        field_type: "text";
        required: true;
    }];

    // Model-predicted sentiment (for comparison; not the Argilla question).
    string predicted_sentiment = 3 [(pgml.argilla_field) = {
        slot: METADATA;
    }];

    // Corrected sentiment after review: "negative" or "positive".
    // When logging a draft from a prediction, set this to the predicted label
    // so Argilla records it as a Suggestion.
    string sentiment = 4 [(pgml.argilla_field) = {
        slot: QUESTION;
        question_type: "label";
        labels: ["negative", "positive"];
        required: true;
    }];

    // Provenance: "model" for suggestion drafts, "human" after correction.
    string source = 5 [(pgml.argilla_field) = {
        slot: METADATA;
    }];
}

// LitServe / client settings for SentimentClassifier.
message SentimentServeConfig {
    option (pgml.kind) = RUN_CONFIG;
    option (pgml.litserve_config) = {
        enable: true;
        service: "SentimentClassifier";
    };

    string url = 1 [(pgml.default) = {string: "http://localhost:8000"}];
    float timeout_s = 2 [(pgml.default) = {float: 30}];
    int32 workers_per_device = 3 [(pgml.default) = {int32: 1}];
    string accelerator = 4 [(pgml.default) = {string: "cpu"}];
}

// Online-loop metrics (predictions vs human feedback).
message SentimentOnlineMetrics {
    option (pgml.kind) = METRIC_SET;
    option (pgml.mlflow) = {
        enable: true;
    };

    int32 n_predictions = 1;
    int32 n_feedback = 2;
    float agreement_rate = 3;
}

// Unary sentiment inference service (LitServe).
service SentimentClassifier {
    option (pgml.litserve) = {
        enable: true;
    };
    rpc Predict(SentimentPredictRequest) returns (SentimentPrediction) {
        option (pgml.litserve_method) = {
            api_path: "/predict";
        };
    }
}

Option	Where	Meaning
(pgml.argilla).enable	message	Required to emit Argilla helpers.
(pgml.argilla).dataset_name	message	Hint returned by *_dataset_name() (default: snake_case message name).
(pgml.argilla_field).*	field	Slot + type/labels (see above).
(pgml.mapper_config).enable	message	Emit *_ALIASES / *_to_row_dict / *_from_row_dict.

Generated module

For SentimentExample, enabling the generator writes *_argilla.py with:

• SentimentExample Pydantic model
• sentiment_example_dataset_name() — e.g. "imdb_sentiment"
• build_sentiment_example_settings(*, client=None)
• to_sentiment_example_record / from_sentiment_example_record
• Optional mapper aliases when configured

Offline Settings construction

Argilla 2.x field/question constructors may try to resolve a default client. Pass an explicit client (or a test double) to avoid requiring API credentials at Settings build time:

from unittest.mock import MagicMock
from pgml_out.sentiment_demo_argilla import build_sentiment_example_settings

settings = build_sentiment_example_settings(client=MagicMock())
assert {f.name for f in settings.fields} == {"text"}
assert {q.name for q in settings.questions} == {"sentiment"}

Records and suggestions

to_*_record places FIELD values in record.fields, METADATA in record.metadata, and non-null QUESTION values as Suggestions (model or weak labels for annotators to confirm). After humans respond in the Argilla UI, map responses back with from_*_record (best-effort for suggestions) or your own response parsing before training.

Suggested flow for classification:

1. Synthesize or load seeds → full models
2. to_*_record → log to Argilla (labels as suggestions)
3. Humans submit responses
4. Export / fetch records → prefer responses over suggestions for train labels
5. Validate back into the same Pydantic / Lance models

Pushing to a live Argilla server

import argilla as rg
from pgml_out.sentiment_demo_argilla import (
    build_sentiment_example_settings,
    sentiment_example_dataset_name,
    to_sentiment_example_record,
)

client = rg.Argilla()  # uses ARGILLA_API_URL / ARGILLA_API_KEY
settings = build_sentiment_example_settings(client=client)
dataset = rg.Dataset(
    name=sentiment_example_dataset_name(),
    settings=settings,
    client=client,
)
dataset.create()
dataset.records.log([to_sentiment_example_record(row) for row in rows])

Or via the bridge:

from py_gen_ml.bridges import synthetic_rows_to_argilla_records

records = synthetic_rows_to_argilla_records(rows, to_record=to_sentiment_example_record)
dataset.records.log(records)

Keep dataset create/log in your module. Do not edit the generated *_argilla.py.

HITL tips

• Put free-text the annotator must read in FIELD, never only in metadata.
• Keep class names in labels: stable—changing them mid-project breaks the UI.
• Use METADATA for provenance (source=imdb|synthetic) so you can filter training sets later.
• For end-to-end sentiment + training wiring, see the Sentiment flywheel example.

See also

• Sentiment flywheel
• PydanticAI synthesis
• Cross-tool bridges
• Message kinds