Sentiment flywheel (IMDB → synthesize → HITL → train → online)

This example project builds a binary sentiment dataset for movie reviews and wires it through the full py-gen-ml loop:

Offline (train-time)

1. Seed with IMDB-style reviews (bundled JSON)
2. Synthesize more labeled rows with PydanticAI (NativeOutput) against OpenAI
3. HITL — push rows to Argilla (FIELD text / QUESTION sentiment) for review
4. Persist accepted rows to LanceDB (optional)
5. Train a small TF–IDF + logistic regression classifier
6. Track the run with MLflow

Online (serve-time)

1. Serve scored reviews via LitServe (SentimentPredictRequest → SentimentPrediction)
2. Store predictions in LanceDB (separate table from labeled examples)
3. Re-label via Argilla on SentimentFeedback (predicted label as suggestion)
4. Store feedback rows and track agreement with MLflow

One protobuf file defines the contracts: labeled SentimentExample, serve I/O, separate SentimentPrediction / SentimentFeedback, plus train and online metrics. Enable several generators and you get adapters together—no duplicated schemas per tool.

flowchart LR
  proto["sentiment_demo.proto"]
  proto --> row["SentimentExample FEATURE_ROW"]
  proto --> pred["SentimentPrediction PREDICTION"]
  proto --> fb["SentimentFeedback FEEDBACK"]
  proto --> cfg["SentimentTrainConfig RUN_CONFIG"]
  proto --> met["SentimentMetrics METRIC_SET"]
  proto --> onlineMet["SentimentOnlineMetrics METRIC_SET"]
  proto --> serve["SentimentClassifier LitServe"]
  row --> pai["*_pydantic_ai.py"]
  row --> arg["*_argilla.py"]
  row --> lance["*_lancedb.py"]
  pred --> lance
  fb --> arg
  fb --> lance
  cfg --> mlflow["*_mlflow.py"]
  cfg --> cli["*_cli_args.py"]
  met --> mlflow
  onlineMet --> mlflow
  serve --> litserve["*_litserve.py"]
  pai --> run["sentiment_flywheel_demo.py"]
  arg --> run
  lance --> run
  mlflow --> run
  cli --> run
  litserve --> online["sentiment_online_demo.py"]
  lance --> online
  arg --> online
  mlflow --> online

flowchart TB
  imdb["IMDB seed JSON"] --> synth["synthesize_sentiment_example"]
  synth --> argilla["Argilla FIELD text / QUESTION sentiment"]
  argilla --> labels["Accepted labels"]
  labels --> lanceStore["LanceDB examples"]
  labels --> train["TF-IDF + LogReg"]
  train --> track["MLflow"]

flowchart TB
  req["SentimentPredictRequest"] --> lit["LitServe Predict"]
  lit --> prediction["SentimentPrediction"]
  prediction --> lancePred["LanceDB predictions"]
  prediction --> fbDraft["SentimentFeedback draft"]
  fbDraft --> argFb["Argilla re-label"]
  argFb --> feedback["SentimentFeedback human"]
  feedback --> lanceFb["LanceDB feedback"]
  prediction --> onlineTrack["MLflow agreement"]
  feedback --> onlineTrack

One schema file, many generators

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";
        };
    }
}

Generated module	What you use it for
sentiment_demo_base.py	Canonical Pydantic / YAML models
sentiment_demo_cli_args.py	Generated Typer CLI flags for SentimentTrainConfig
sentiment_demo_pydantic_ai.py	Full + Partial models, synthesize_sentiment_example(_sync)
sentiment_demo_argilla.py	Argilla Settings for examples and feedback
sentiment_demo_lancedb.py	LanceModels for examples, predictions, feedback
sentiment_demo_litserve.py	create_sentiment_classifier_server, predict client
sentiment_demo_mlflow.py	Train + online metric helpers

SentimentExample field layout (offline HITL):

Field	Use	Argilla slot
id	Stable row id	METADATA
source	imdb or synthetic	METADATA
text	Review body (classifier input)	FIELD
sentiment	negative / positive	QUESTION

Online contracts (separate messages — do not overload SentimentExample):

Message	Kind	Notable fields
SentimentPredictRequest	FEATURE_ROW	id, text (RPC input)
SentimentPrediction	PREDICTION	sample_id, text, sentiment, score, model_version
SentimentFeedback	FEEDBACK	sample_id, text (FIELD), predicted_sentiment (METADATA), sentiment (QUESTION), source
SentimentServeConfig	RUN_CONFIG	LitServe URL / workers / accelerator
SentimentOnlineMetrics	METRIC_SET	n_predictions, n_feedback, agreement_rate

Tracking messages use shared (pgml.tracking_field) slots (see MLflow):

Message	Kind	Notable fields
SentimentTrainConfig	RUN_CONFIG	run_name (TAG), synthesize_count, diversify_rounds, openai_model, test_size (PARAM)
SentimentMetrics	METRIC_SET	accuracy, n_train, n_test, n_labeled

Proto leading comments on text / sentiment become Field(description=...) and thus JSON Schema descriptions for synthesis. Keep them specific—vague comments produce weaker synthetic data (see PydanticAI).

Enable the generators when you regenerate:

py-gen-ml path/to/sentiment_demo.proto \
  --generators=base,patch,sweep,cli_args,pydantic_ai,argilla,lancedb,litserve,mlflow

Setup

From the docs snippets project:

cd docs/snippets
uv sync --extra bridges --extra pydantic-ai --extra argilla --extra lancedb --extra litserve --extra mlflow
bash regenerate.sh   # regenerates all snippet protos, including sentiment_demo

OpenAI credentials (required)

The demo talks to OpenAI (Azure OpenAI–compatible endpoint) via PydanticAI. Export:

export OPENAI_API_KEY=...
export OPENAI_ENDPOINT=https://YOUR_RESOURCE.openai.azure.com/
export OPENAI_API_VERSION=2024-12-01-preview

Variable	Meaning
OPENAI_API_KEY	API key
OPENAI_ENDPOINT	Azure OpenAI resource endpoint (or compatible base URL)
OPENAI_API_VERSION	Azure API version (e.g. 2024-12-01-preview)

The deployment / model name is SentimentTrainConfig.openai_model (YAML default gpt-4o, overridable with --openai-model).

Argilla credentials (required to push HITL records)

export ARGILLA_API_URL=https://...
export ARGILLA_API_KEY=...

MLflow uses its normal env / local tracking URI (no cloud account required for a file store).

Seed file (short IMDB-style reviews checked into the repo):

docs/snippets/data/imdb_sentiment_seeds.json

These are not the full Stanford IMDB corpus—they are compact, realistic reviews shaped like IMDB user text so the demo stays fast and license-friendly. Swap in your own seeds (same JSON shape) without changing the proto.

Run the demo

SentimentTrainConfig is populated like CIFAR: YAML base config + generated CLI overrides via @pgml.pgml_cmd / apply_cli_args.

Default YAML:

# yaml-language-server: $schema=schemas/sentiment_train_config.json
run_name: sentiment
synthesize_count: 4
diversify_rounds: 1
openai_model: gpt-4o
test_size: 0.25

"""Sentiment flywheel demo: IMDB seeds → synthesize → Argilla HITL → train + MLflow.

Run from ``docs/snippets`` after codegen::

    uv sync --extra bridges --extra pydantic-ai --extra argilla --extra lancedb --extra mlflow

    export OPENAI_API_KEY=...
    export OPENAI_ENDPOINT=https://YOUR_RESOURCE.openai.azure.com/
    export OPENAI_API_VERSION=2024-12-01-preview

    # optional: log records to a live Argilla server
    # export ARGILLA_API_URL=...
    # export ARGILLA_API_KEY=...

    uv run python -m snippets.sentiment_flywheel_demo \\
      --config-paths configs/base/sentiment_train_config.yaml

    # override fields via generated CLI flags (same as CIFAR):
    #   --synthesize-count 2 --openai-model gpt-4o --run-name my-run
"""
import json
import os
import tempfile
from pathlib import Path
from typing import List, Optional, Sequence

import argilla as rg
import pgml_out.sentiment_demo_base as base
import pgml_out.sentiment_demo_cli_args as cli_args
import typer
from pgml_out.sentiment_demo_argilla import (
    build_sentiment_example_settings,
    sentiment_example_dataset_name,
    to_sentiment_example_record,
)
from pgml_out.sentiment_demo_lancedb import (
    SentimentExample as LanceSentimentExample,
    create_sentiment_example_table,
    sentiment_example_merge_on,
    sentiment_example_table_name,
)
from pgml_out.sentiment_demo_pydantic_ai import (
    SentimentExample,
    SentimentExamplePartial,
    synthesize_sentiment_example_sync,
)
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.azure import AzureProvider

import py_gen_ml as pgml
from py_gen_ml.bridges.lancedb_rows import merge_rows
from py_gen_ml.bridges.synthesis_argilla import synthetic_rows_to_argilla_records

# .../docs/snippets/src/snippets/this_file.py → parents[2] == docs/snippets
SNIPPETS_ROOT = Path(__file__).resolve().parents[2]
SEEDS_PATH = SNIPPETS_ROOT / 'data' / 'imdb_sentiment_seeds.json'
DEFAULT_CONFIG_PATH = SNIPPETS_ROOT / 'configs' / 'base' / 'sentiment_train_config.yaml'

SYSTEM_PROMPT = """\
You generate realistic IMDB-style movie reviews for binary sentiment classification.
Each example must include:
- text: a short user review (1-4 sentences, informal, specific to a fictional film)
- sentiment: exactly "positive" or "negative"
- source: always "synthetic"
- id: a unique string like "synth-001"
Match the tone and length of the few-shot IMDB seeds. Prefer diversity across genres,
tone, and vocabulary. Do not copy seed text verbatim.
"""

app = typer.Typer(pretty_exceptions_enable=False)


def openai_model_from_env(*, model: Optional[str] = None) -> OpenAIChatModel:
    """Build an OpenAI chat model from endpoint, key, API version, and model name.

    ``model`` comes from ``SentimentTrainConfig.openai_model`` when running the
    flywheel; otherwise ``OPENAI_MODEL`` is required.
    """
    return OpenAIChatModel(
        model if model is not None else os.environ['OPENAI_MODEL'],
        provider=AzureProvider(
            azure_endpoint=os.environ['OPENAI_ENDPOINT'],
            api_version=os.environ['OPENAI_API_VERSION'],
            api_key=os.environ['OPENAI_API_KEY'],
        ),
    )


def argilla_client_from_env() -> rg.Argilla:
    """Build an Argilla client from ``ARGILLA_API_URL`` / ``ARGILLA_API_KEY``."""
    return rg.Argilla(
        api_url=os.environ['ARGILLA_API_URL'],
        api_key=os.environ['ARGILLA_API_KEY'],
    )


def load_imdb_seeds(path: Path = SEEDS_PATH) -> List[SentimentExample]:
    """Load bundled IMDB-style seed reviews as full ``SentimentExample`` rows."""
    raw = json.loads(path.read_text())
    return [SentimentExample.model_validate(row) for row in raw]


def push_argilla_dataset(
    *,
    client: rg.Argilla,
    settings: rg.Settings,
    records: Sequence[rg.Record],
) -> rg.Dataset:
    """Create (or reuse) the Argilla dataset and log ``records`` as suggestions."""
    name = sentiment_example_dataset_name()
    dataset = client.datasets(name=name)
    if dataset is None:
        dataset = rg.Dataset(name=name, settings=settings, client=client)
        dataset.create()
    dataset.records.log(list(records))
    return dataset


def train_sentiment_classifier(
    rows: Sequence[SentimentExample],
    *,
    test_size: float = 0.25,
) -> dict:
    """Train a small TF-IDF + logistic regression sentiment model.

    Returns accuracy on a simple holdout split.
    """
    from sklearn.feature_extraction.text import TfidfVectorizer
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import accuracy_score
    from sklearn.model_selection import train_test_split
    from sklearn.pipeline import Pipeline

    texts = [r.text for r in rows]
    labels = [r.sentiment for r in rows]
    x_train, x_test, y_train, y_test = train_test_split(
        texts,
        labels,
        test_size=test_size,
        random_state=42,
        stratify=labels if len(set(labels)) > 1 else None,
    )
    pipe = Pipeline(
        [
            ('tfidf', TfidfVectorizer(ngram_range=(1, 2), min_df=1)),
            ('clf', LogisticRegression(max_iter=1000)),
        ],
    )
    pipe.fit(x_train, y_train)
    pred = pipe.predict(x_test)
    return {
        'n_train': len(x_train),
        'n_test': len(x_test),
        'accuracy': float(accuracy_score(y_test, pred)),
        'pipeline': pipe,
    }


def persist_to_lancedb(rows: Sequence[SentimentExample], db_uri: str) -> str:
    """Write examples to a local LanceDB table; return table name."""
    import lancedb

    db = lancedb.connect(db_uri)
    lance_rows = [LanceSentimentExample.model_validate(r.model_dump()) for r in rows]
    table = create_sentiment_example_table(db, mode='overwrite', exist_ok=False)
    merge_rows(table, lance_rows, on=sentiment_example_merge_on())
    return sentiment_example_table_name()


def log_training_to_mlflow(*, train_config: base.SentimentTrainConfig, metrics) -> None:
    """Log ``SentimentTrainConfig`` + ``SentimentMetrics`` to MLflow."""
    from pgml_out.sentiment_demo_mlflow import (
        SentimentTrainConfig as MlflowTrainConfig,
        log_sentiment_metrics,
        start_sentiment_train_config_run,
    )

    mlflow_config = MlflowTrainConfig.model_validate(train_config.model_dump())
    with start_sentiment_train_config_run(mlflow_config):
        log_sentiment_metrics(metrics)


def run_flywheel(
    train_config: base.SentimentTrainConfig,
    *,
    use_lancedb: bool = True,
    push_to_argilla: bool = True,
) -> dict:
    """End-to-end: seeds → OpenAI synthesis → Argilla HITL → LanceDB → train + MLflow."""
    from pgml_out.sentiment_demo_mlflow import SentimentMetrics

    seeds = load_imdb_seeds()
    synthetic = synthesize_sentiment_example_sync(
        model=openai_model_from_env(model=train_config.openai_model),
        system_prompt=SYSTEM_PROMPT,
        count=train_config.synthesize_count,
        examples=[
            SentimentExamplePartial.model_validate(s.model_dump()) for s in seeds
        ],
        diversify_rounds=train_config.diversify_rounds,
    )
    # Until humans respond in Argilla, train on model-suggested labels.
    # After HITL, replace with labels from Argilla responses.
    labeled = list(seeds) + list(synthetic)

    client = argilla_client_from_env()
    settings = build_sentiment_example_settings(client=client)
    records = synthetic_rows_to_argilla_records(
        labeled,
        to_record=to_sentiment_example_record,
    )
    if push_to_argilla:
        push_argilla_dataset(client=client, settings=settings, records=records)

    result: dict = {
        'n_seeds': len(seeds),
        'n_synthetic': len(synthetic),
        'n_labeled': len(labeled),
        'n_argilla_records': len(records),
        'n_settings_fields': len(settings.fields),
        'n_settings_questions': len(settings.questions),
        'argilla_dataset': sentiment_example_dataset_name(),
        'openai_model': train_config.openai_model,
        'run_name': train_config.run_name,
    }

    if use_lancedb:
        with tempfile.TemporaryDirectory() as tmp:
            table_name = persist_to_lancedb(labeled, tmp)
            result['lancedb_table'] = table_name

    train_metrics = train_sentiment_classifier(
        labeled,
        test_size=train_config.test_size,
    )
    result['train_accuracy'] = train_metrics['accuracy']
    result['n_train'] = train_metrics['n_train']
    result['n_test'] = train_metrics['n_test']

    metrics = SentimentMetrics(
        accuracy=train_metrics['accuracy'],
        n_train=train_metrics['n_train'],
        n_test=train_metrics['n_test'],
        n_labeled=len(labeled),
    )
    log_training_to_mlflow(train_config=train_config, metrics=metrics)
    result['mlflow_experiment'] = 'imdb_sentiment'
    return result


@pgml.pgml_cmd(app=app)
def main(
    config_paths: List[str] = typer.Option(
        default_factory=lambda: [str(DEFAULT_CONFIG_PATH)],
        help='Paths to SentimentTrainConfig YAML files',
    ),
    cli_args: cli_args.SentimentTrainConfigArgs = typer.Option(...),
    use_lancedb: bool = typer.Option(True, help='Write labeled rows to a temp LanceDB'),
    push_to_argilla: bool = typer.Option(True, help='Push records to Argilla'),
) -> None:
    train_config = base.SentimentTrainConfig.from_yaml_files(config_paths)
    train_config = train_config.apply_cli_args(cli_args)
    summary = run_flywheel(
        train_config,
        use_lancedb=use_lancedb,
        push_to_argilla=push_to_argilla,
    )
    print(json.dumps({k: v for k, v in summary.items() if k != 'pipeline'}, indent=2))


if __name__ == '__main__':
    app()

cd docs/snippets
uv run python -m snippets.sentiment_flywheel_demo \
  --config-paths configs/base/sentiment_train_config.yaml

# override any train-config field:
uv run python -m snippets.sentiment_flywheel_demo \
  --config-paths configs/base/sentiment_train_config.yaml \
  --synthesize-count 2 --run-name my-run

Example output shape:

{
  "n_seeds": 8,
  "n_synthetic": 8,
  "n_labeled": 16,
  "n_argilla_records": 16,
  "n_settings_fields": 1,
  "n_settings_questions": 1,
  "argilla_dataset": "imdb_sentiment",
  "openai_model": "gpt-4o",
  "run_name": "sentiment",
  "lancedb_table": "sentiment_examples",
  "train_accuracy": 0.75,
  "n_train": 12,
  "n_test": 4,
  "mlflow_experiment": "imdb_sentiment"
}

Accuracy on this tiny toy set is not meaningful; the point is the wiring from seed → synth → HITL records → store → train → track.

Step-by-step

1. Load IMDB seeds

from snippets.sentiment_flywheel_demo import load_imdb_seeds

seeds = load_imdb_seeds()
assert seeds[0].source == "imdb"
assert seeds[0].sentiment in {"positive", "negative"}

Seeds validate as generated SentimentExample (full) models.

2. Synthesize with few-shot + diversify

from snippets.sentiment_flywheel_demo import SYSTEM_PROMPT, openai_model_from_env
from pgml_out.sentiment_demo_base import SentimentTrainConfig
from pgml_out.sentiment_demo_pydantic_ai import (
    SentimentExamplePartial,
    synthesize_sentiment_example_sync,
)

train_config = SentimentTrainConfig.from_yaml_files(
    ["configs/base/sentiment_train_config.yaml"]
)
synthetic = synthesize_sentiment_example_sync(
    model=openai_model_from_env(model=train_config.openai_model),
    system_prompt=SYSTEM_PROMPT,
    count=train_config.synthesize_count,
    examples=[SentimentExamplePartial.model_validate(s.model_dump()) for s in seeds],
    diversify_rounds=train_config.diversify_rounds,
)

IMDB seeds are passed as few-shot examples. The system prompt asks for IMDB-style reviews with source="synthetic".

diversify_rounds=1 means: generate a first batch of count rows, then run another round that feeds prior full outputs back as examples. Total synthetic rows ≈ count * (diversify_rounds + 1).

3. HITL via Argilla

from snippets.sentiment_flywheel_demo import argilla_client_from_env, push_argilla_dataset
from pgml_out.sentiment_demo_argilla import (
    build_sentiment_example_settings,
    to_sentiment_example_record,
)
from py_gen_ml.bridges import synthetic_rows_to_argilla_records

client = argilla_client_from_env()
settings = build_sentiment_example_settings(client=client)
records = synthetic_rows_to_argilla_records(
    seeds + synthetic,
    to_record=to_sentiment_example_record,
)
dataset = push_argilla_dataset(client=client, settings=settings, records=records)
assert {f.name for f in settings.fields} == {"text"}
assert {q.name for q in settings.questions} == {"sentiment"}

Annotators confirm or correct sentiment in the Argilla UI. Labels are logged as suggestions first; prefer human responses for training once review is done.

4. Store in LanceDB

run_flywheel(train_config, use_lancedb=True) writes labeled rows to a temp LanceDB using the generated SentimentExample LanceModel and py_gen_ml.bridges.merge_rows.

In production, keep a durable URI and append only after HITL acceptance.

5. Train

train_sentiment_classifier fits TfidfVectorizer + LogisticRegression and reports holdout accuracy. Swap this for your real training stack; the dataset contract stays the same protobuf message.

6. Track with MLflow

run_flywheel logs the same SentimentTrainConfig instance that drove synthesis and training (from YAML + CLI), plus SentimentMetrics:

from snippets.sentiment_flywheel_demo import log_training_to_mlflow
from pgml_out.sentiment_demo_base import SentimentTrainConfig
from pgml_out.sentiment_demo_mlflow import SentimentMetrics

config = SentimentTrainConfig.from_yaml_files(
    ["configs/base/sentiment_train_config.yaml"]
)
metrics = SentimentMetrics(accuracy=0.75, n_train=12, n_test=4, n_labeled=16)
log_training_to_mlflow(train_config=config, metrics=metrics)

Typical next steps after offline training:

• Load accepted rows from LanceDB (load_seeds_from_table) or Argilla exports
• Filter source / quality metadata
• Train / evaluate / push a serving model (see Online loop below)

Online loop

After you have a classifier, the same proto drives serve → prediction → feedback.

cd docs/snippets
uv run python -m snippets.sentiment_online_demo score-batch
# live server (writes predictions to ./sentiment_online.lancedb):
uv run python -m snippets.sentiment_online_demo serve --port 8000

"""Sentiment online loop: feature → LitServe → prediction → feedback → store / track.

Run from ``docs/snippets`` after codegen::

    uv sync --extra bridges --extra lancedb --extra argilla --extra litserve --extra mlflow

    # optional live Argilla for HITL push
    # export ARGILLA_API_URL=...
    # export ARGILLA_API_KEY=...

    # CI-friendly batch (no long-running server):
    uv run python -m snippets.sentiment_online_demo score-batch

    # Live LitServe:
    uv run python -m snippets.sentiment_online_demo serve
"""
from __future__ import annotations

import json
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import typer

from snippets.sentiment_flywheel_demo import (
    SEEDS_PATH,
    argilla_client_from_env,
    load_imdb_seeds,
    train_sentiment_classifier,
)

MODEL_VERSION = 'tfidf-logreg-v1'

app = typer.Typer(pretty_exceptions_enable=False)


def fit_pipeline_from_seeds(path: Path = SEEDS_PATH) -> Any:
    """Fit TF–IDF + LogReg on bundled IMDB seeds; return the sklearn pipeline."""
    seeds = load_imdb_seeds(path)
    metrics = train_sentiment_classifier(seeds, test_size=0.25)
    return metrics['pipeline']


def predict_sentiment(
    pipeline: Any,
    *,
    sample_id: str,
    text: str,
    model_version: str = MODEL_VERSION,
) -> Any:
    """Score one review; return a LitServe ``SentimentPrediction``."""
    from pgml_out.sentiment_demo_litserve import SentimentPrediction

    label = str(pipeline.predict([text])[0])
    proba = pipeline.predict_proba([text])[0]
    classes = list(pipeline.classes_)
    score = float(proba[classes.index(label)])
    return SentimentPrediction(
        sample_id=sample_id,
        text=text,
        sentiment=label,
        score=score,
        model_version=model_version,
    )


def merge_to_feedback(_request: Any, prediction: Any) -> Any:
    """Build a ``SentimentFeedback`` draft (predicted label as Argilla suggestion)."""
    from pgml_out.sentiment_demo_argilla import SentimentFeedback

    return SentimentFeedback(
        sample_id=prediction.sample_id,
        text=prediction.text,
        predicted_sentiment=prediction.sentiment,
        sentiment=prediction.sentiment,
        source='model',
    )


def apply_human_corrections(
    drafts: Sequence[Any],
    *,
    ground_truth: Optional[Dict[str, str]] = None,
) -> List[Any]:
    """Simulate HITL: set ``source=human`` and optionally override ``sentiment``.

    When ``ground_truth`` maps ``sample_id → label``, use that as the corrected
    label (demo / tests). Otherwise keep the predicted suggestion as the response.
    """
    from pgml_out.sentiment_demo_argilla import SentimentFeedback

    out: List[Any] = []
    for draft in drafts:
        corrected = (
            ground_truth.get(draft.sample_id, draft.sentiment)
            if ground_truth is not None
            else draft.sentiment
        )
        out.append(
            SentimentFeedback(
                sample_id=draft.sample_id,
                text=draft.text,
                predicted_sentiment=draft.predicted_sentiment,
                sentiment=corrected,
                source='human',
            ),
        )
    return out


def ensure_tables(db: Any) -> Tuple[Any, Any]:
    """Create (or open) prediction and feedback LanceDB tables."""
    from pgml_out.sentiment_demo_lancedb import (
        create_sentiment_feedback_table,
        create_sentiment_prediction_table,
    )

    return (
        create_sentiment_prediction_table(db, mode='overwrite', exist_ok=False),
        create_sentiment_feedback_table(db, mode='overwrite', exist_ok=False),
    )


def build_feedback_records(
    requests: Sequence[Any],
    predictions: Sequence[Any],
) -> Tuple[List[Any], List[Any]]:
    """Merge each request/prediction into a feedback draft and Argilla record."""
    from pgml_out.sentiment_demo_argilla import to_sentiment_feedback_record
    from py_gen_ml.bridges import log_prediction_for_review

    drafts: List[Any] = []
    records: List[Any] = []
    for req, pred in zip(requests, predictions):
        draft = merge_to_feedback(req, pred)
        drafts.append(draft)
        records.append(
            log_prediction_for_review(
                req,
                pred,
                merge=merge_to_feedback,
                to_record=to_sentiment_feedback_record,
            ),
        )
    return drafts, records


def push_feedback_argilla(*, records: Sequence[Any]) -> str:
    """Create (or reuse) the feedback dataset and log ``records``."""
    import argilla as rg
    from pgml_out.sentiment_demo_argilla import (
        build_sentiment_feedback_settings,
        sentiment_feedback_dataset_name,
    )

    name = sentiment_feedback_dataset_name()
    client = argilla_client_from_env()
    settings = build_sentiment_feedback_settings(client=client)
    dataset = client.datasets(name=name)
    if dataset is None:
        dataset = rg.Dataset(name=name, settings=settings, client=client)
        dataset.create()
    dataset.records.log(list(records))
    return name


def log_online_metrics_to_mlflow(
    *,
    n_predictions: int,
    n_feedback: int,
    agreement: float,
) -> None:
    """Log ``SentimentOnlineMetrics`` under a short-lived MLflow run."""
    import mlflow
    from pgml_out.sentiment_demo_mlflow import (
        SentimentOnlineMetrics,
        log_sentiment_online_metrics,
    )

    metrics = SentimentOnlineMetrics(
        n_predictions=n_predictions,
        n_feedback=n_feedback,
        agreement_rate=agreement,
    )
    mlflow.set_experiment('imdb_sentiment_online')
    with mlflow.start_run(run_name='online-batch'):
        log_sentiment_online_metrics(metrics)


def score_batch(
    *,
    db_uri: Optional[str] = None,
    push_to_argilla: bool = False,
    use_ground_truth: bool = True,
) -> dict:
    """Offline online-loop: fit → predict seeds → store → feedback → track."""
    import lancedb
    from pgml_out.sentiment_demo_lancedb import (
        SentimentFeedback as LanceFeedback,
        SentimentPrediction as LancePrediction,
        sentiment_feedback_merge_on,
        sentiment_prediction_merge_on,
    )
    from pgml_out.sentiment_demo_litserve import SentimentPredictRequest
    from pgml_out.sentiment_demo_argilla import sentiment_feedback_dataset_name
    from py_gen_ml.bridges import merge_rows

    seeds = load_imdb_seeds()
    pipeline = fit_pipeline_from_seeds()

    requests = [
        SentimentPredictRequest(id=s.id, text=s.text) for s in seeds
    ]
    predictions = [
        predict_sentiment(pipeline, sample_id=req.id, text=req.text) for req in requests
    ]
    drafts, records = build_feedback_records(requests, predictions)
    ground_truth = {s.id: s.sentiment for s in seeds} if use_ground_truth else None
    feedbacks = apply_human_corrections(drafts, ground_truth=ground_truth)

    dataset_name = sentiment_feedback_dataset_name()
    if push_to_argilla:
        dataset_name = push_feedback_argilla(records=records)

    owns_tmp = db_uri is None
    tmp_ctx = tempfile.TemporaryDirectory() if owns_tmp else None
    try:
        uri = db_uri if db_uri is not None else tmp_ctx.name  # type: ignore[union-attr]
        db = lancedb.connect(uri)
        pred_table, fb_table = ensure_tables(db)
        merge_rows(
            pred_table,
            [LancePrediction.model_validate(p.model_dump()) for p in predictions],
            on=sentiment_prediction_merge_on(),
        )
        merge_rows(
            fb_table,
            [LanceFeedback.model_validate(f.model_dump()) for f in feedbacks],
            on=sentiment_feedback_merge_on(),
        )
        pred_by_id = {p.sample_id: p.sentiment for p in predictions}
        matched = sum(1 for f in feedbacks if pred_by_id.get(f.sample_id) == f.sentiment)
        agreement = matched / len(feedbacks) if feedbacks else 0.0
        log_online_metrics_to_mlflow(
            n_predictions=len(predictions),
            n_feedback=len(feedbacks),
            agreement=agreement,
        )
        return {
            'n_predictions': len(predictions),
            'n_feedback': len(feedbacks),
            'n_argilla_records': len(records),
            'argilla_dataset': dataset_name,
            'agreement_rate': agreement,
            'lancedb_predictions': 'sentiment_predictions',
            'lancedb_feedback': 'sentiment_feedback',
            'mlflow_experiment': 'imdb_sentiment_online',
            'model_version': MODEL_VERSION,
            'db_uri': uri if not owns_tmp else None,
        }
    finally:
        if tmp_ctx is not None:
            tmp_ctx.cleanup()


@app.command('score-batch')
def score_batch_cmd(
    push_to_argilla: bool = typer.Option(False, help='Push feedback drafts to Argilla'),
    db_uri: Optional[str] = typer.Option(None, help='LanceDB URI (temp dir if omitted)'),
) -> None:
    """Run the online loop offline over seed texts (CI-friendly)."""
    summary = score_batch(db_uri=db_uri, push_to_argilla=push_to_argilla)
    print(json.dumps(summary, indent=2))


@app.command('serve')
def serve_cmd(
    port: int = typer.Option(8000, help='LitServe port'),
    db_uri: str = typer.Option('./sentiment_online.lancedb', help='LanceDB URI'),
) -> None:
    """Fit on seeds and serve ``SentimentClassifier`` via LitServe."""
    import lancedb
    from pgml_out.sentiment_demo_base import SentimentServeConfig
    from pgml_out.sentiment_demo_lancedb import (
        SentimentPrediction as LancePrediction,
        sentiment_prediction_merge_on,
    )
    from pgml_out.sentiment_demo_litserve import (
        SentimentPredictRequest,
        create_sentiment_classifier_server,
    )
    from py_gen_ml.bridges import merge_rows

    pipeline = fit_pipeline_from_seeds()
    db = lancedb.connect(db_uri)
    pred_table, _fb_table = ensure_tables(db)

    def predict(request: SentimentPredictRequest):
        prediction = predict_sentiment(
            pipeline,
            sample_id=request.id,
            text=request.text,
        )
        merge_rows(
            pred_table,
            [LancePrediction.model_validate(prediction.model_dump())],
            on=sentiment_prediction_merge_on(),
        )
        return prediction

    server = create_sentiment_classifier_server(
        predict=predict,
        config=SentimentServeConfig(
            accelerator='cpu',
            workers_per_device=1,
            timeout_s=30.0,
            url=f'http://127.0.0.1:{port}',
        ),
    )
    server.run(host='127.0.0.1', port=port, generate_client_file=False)


if __name__ == '__main__':
    app()

Flow:

1. Fit TF–IDF + LogReg on IMDB seeds (same helper as the offline demo).
2. Score each seed as SentimentPredictRequest → SentimentPrediction.
3. Persist predictions with merge_rows to the sentiment_predictions table.
4. Merge request + prediction into a SentimentFeedback draft (source=model); Argilla gets the predicted label as a Suggestion on the QUESTION field.
5. Simulate (or apply) human corrections → source=human feedback rows in sentiment_feedback.
6. Log SentimentOnlineMetrics (n_predictions, n_feedback, agreement_rate) to MLflow experiment imdb_sentiment_online.

SentimentPrediction and SentimentFeedback are separate messages from SentimentExample. Do not overload the labeled training row for prod scores.

See also LitServe and bridges.serving_argilla.

Completing partial reviews

If you have review text without labels (or the reverse), use Path B (incomplete):

from snippets.sentiment_flywheel_demo import openai_model_from_env
from pgml_out.sentiment_demo_pydantic_ai import (
    SentimentExamplePartial,
    synthesize_sentiment_example_sync,
)

filled = synthesize_sentiment_example_sync(
    model=openai_model_from_env(model="gpt-4o"),
    system_prompt="Fill only missing fields for sentiment examples.",
    incomplete=[
        SentimentExamplePartial(
            id="need-label-1",
            source="imdb",
            text="An absolute masterpiece of modern cinema.",
            sentiment=None,
        ),
    ],
)

Only sentiment is sent in the gap JSON Schema; provided fields are preserved. Do not combine incomplete with diversify_rounds > 0 in one call (v1).

Production checklist

• Field comments on every synthesis-relevant field
• Explicit Argilla slot on every field (FIELD / QUESTION / METADATA)
• Seed set balanced enough for few-shot (both classes)
• OPENAI_API_KEY / OPENAI_ENDPOINT / OPENAI_API_VERSION set (offline synth)
• HITL: train on human responses, not only LLM suggestions
• Track source / id in metadata for auditability
• Keep predictions and feedback in separate LanceDB tables from labeled examples
• MLflow tracking URI configured for real runs

See also

• PydanticAI synthesis
• Argilla datasets
• MLflow tracking
• LitServe services
• Cross-tool bridges
• LanceDB schemas
• Message kinds