sentimentizer

Lightweight PyTorch models for sentiment analysis. Small models can be pretty effective for classification tasks at a much smaller cost to deploy — all models were trained on a single GPU in minutes, and inference requires less than 1GB of memory.

Beta release — API is subject to change.

Install

# Install local-only version (no Ray dependency)
uv add sentimentizer

# Install with distributed training, tuning, and serving features
uv add "sentimentizer[ray]"

# Install the headless ComfyUI image API client
uv add "sentimentizer[diffusion]"

# Install all features
uv add "sentimentizer[ray,diffusion]"

Quick Start

Run a pre-trained model locally:

from sentimentizer.predictor import SentimentPredictor

# Load the model
predictor = SentimentPredictor(model_name="encoder")

# Predict sentiment (returns label, score, all-class scores, token count, and model type)
result = predictor.predict("amazing restaurant!")
# >> {"positive": 0.848753035068512,
#     "label": "positive",
#     "score": 0.848753035068512,
#     "scores": {"negative": 0.05746544152498245,
#                "neutral": 0.09378153085708618,
#                "positive": 0.848753035068512},
#     "token_count": 2,
#     "model": "encoder"}

# Batch prediction
results = predictor.predict_batch(["Great food!", "Terrible service."])
# >> [{'positive': 0.848753035068512,
#      'label': 'positive',
#      'score': 0.848753035068512,
#      'scores': {'negative': 0.05746544152498245,
#                 'neutral': 0.09378153085708618,
#                 'positive': 0.848753035068512},
#      'token_count': 2,
#      'model': 'encoder'},
#     {'negative': 0.9586858749389648,
#      'label': 'negative',
#      'score': 0.9586858749389648,
#      'scores': {'negative': 0.9586858749389648,
#                 'neutral': 0.03224344924092293,
#                 'positive': 0.0090706842020154},
#      'token_count': 2,
#      'model': 'encoder'}]

Models output 3-class probabilities (negative, neutral, positive) that sum to 1.0 per sample.

Image Generation (Krea 2 / Ideogram 4)

Image generation uses a separately managed, headless ComfyUI process. Torch Sentiment submits native Krea 2 and Ideogram 4 workflows over HTTP; it does not load image weights or reserve the CUDA device inside Ray. The supported checkpoints are the official INT8 ConvRot builds.

Prerequisites

# Torch Sentiment side
uv sync --extra ray --extra diffusion

# Separate ComfyUI checkout and environment. This revision has native ConvRot,
# Krea 2, and Ideogram 4 nodes; no custom nodes or browser UI are required.
git clone https://github.com/Comfy-Org/ComfyUI.git ../ComfyUI
cd ../ComfyUI
git checkout feca51a
uv venv
uv pip install -r requirements.txt

# Start API-only on localhost. Install the CUDA PyTorch build appropriate for
# the host before this step if requirements did not select it.
uv run python main.py --listen 127.0.0.1 --port 8188 --disable-auto-launch \
  --temp-directory /tmp/sentimentizer-comfyui

Configure

Put these files under the ComfyUI checkout:

The official download URLs are listed in docs/serving.md.

Configuration lives in two YAML files:

• sentimentizer/serve/service.yaml — enabled models, ComfyUI URL, auth, and rate limits.
• sentimentizer/diffusion/diffusion_config.yaml — checkpoint filenames and generation defaults.

Edit service.yaml to enable a model:

# Enable one or both models
krea_2_enabled: true
# Required for Krea 2. Review the Community License and confirm eligibility.
krea_2_license_accepted: true
ideogram_4_enabled: true
# Required for Ideogram 4. Review its non-commercial model license first.
ideogram_4_license_accepted: true
# Required for Krea 2. The endpoint must explicitly approve generated output.
image_moderation_url: "http://127.0.0.1:8090/v1/moderate-image"
image_moderation_api_key: "replace-me"
default_image_model: "krea_2"
comfyui_base_url: "http://127.0.0.1:8188"
comfyui_temp_directory: "/tmp/sentimentizer-comfyui/temp"

# Auth — required for image routes (/v1/images/*)
api_keys: ["sk-your-secret-key"]

Or via environment variables:

export SENTIMENTIZER_KREA_2_ENABLED=true
export SENTIMENTIZER_KREA_2_LICENSE_ACCEPTED=true
export SENTIMENTIZER_IDEOGRAM_4_ENABLED=true
export SENTIMENTIZER_IDEOGRAM_4_LICENSE_ACCEPTED=true
export SENTIMENTIZER_IMAGE_MODERATION_URL=http://127.0.0.1:8090/v1/moderate-image
export SENTIMENTIZER_COMFYUI_BASE_URL=http://127.0.0.1:8188
export SENTIMENTIZER_COMFYUI_TEMP_DIRECTORY=/tmp/sentimentizer-comfyui/temp
export SENTIMENTIZER_API_KEYS=sk-your-secret-key

Krea 2 uses the Krea 2 Community License. Its output-moderation endpoint receives JSON containing image_b64, mime_type, model, and user, and must return {"safe": true} to release the image. A rejection may return {"safe": false, "code": "...", "message": "..."}. Missing, malformed, or unavailable moderation responses fail closed.

Run

# Start Ray Serve; ComfyUI must already be running.
python -m sentimentizer.serve

# Krea 2 generation
curl -X POST http://localhost:8000/v1/images/generate \
  -H "Authorization: Bearer sk-your-secret-key" \
  -H "Content-Type: application/json" \
  -d '{"prompt": "a cinematic portrait of an astronaut", "model": "krea_2", "width": 1024, "height": 1024}'

# Ideogram 4 generation; structured JSON prompts give its strongest text/layout results.
curl -X POST http://localhost:8000/v1/images/generate \
  -H "Authorization: Bearer sk-your-secret-key" \
  -H "Content-Type: application/json" \
  -d '{"prompt": "a bold poster reading FRESH", "model": "ideogram_4", "width": 1024, "height": 1024}'

# List available models
curl http://localhost:8000/v1/images/models \
  -H "Authorization: Bearer sk-your-secret-key"

# Async job mode (for long-running requests)
curl -X POST http://localhost:8000/v1/images/jobs \
  -H "Authorization: Bearer sk-your-secret-key" \
  -H "Content-Type: application/json" \
  -d '{"prompt": "a cinematic portrait of an astronaut", "model": "krea_2"}'

# Poll job status
curl http://localhost:8000/v1/images/jobs/{job_id} \
  -H "Authorization: Bearer sk-your-secret-key"

API Endpoints

Models

Four architectures are available:

All models output 3-class logits (B, 3) mapped to: negative (0), neutral (1), positive (2).

Documentation

Detailed guides and implementation details are available in the specialized documentation files:

• 🚀 Model Serving Guide: Ray Serve application deployment, FastAPI endpoints (sentiment/routing/image generation), and the Go CLI client.
• 🎨 Headless Image Serving: Krea 2 and Ideogram 4 setup, checkpoints, API, and licensing.
• 🏋️ Model Training & Checkpointing Guide: Yelp datasets, single-node/distributed commands, training arguments, sleep prevention, and checkpoint resuming.
• ⚙️ Model Configuration Reference: Configuration dataclasses (RNNConfig, EncoderConfig, etc.), parameter defaults, and consistency checks.
• 🎛️ Hyperparameter Tuning Guide: Optuna searches, LangGraph iterative agent tuning (via Ollama GLM 5.1), and validation/retries.
• 🔗 Hugging Face Hub Integration: Pre-trained weights synchronization, explicit pull/push, and auto-generated model cards.
• 📈 Metrics and Monitoring Pipeline: Exporter details, Grafana dashboards, Prometheus scrape targets, NaN handling, and real-time intra-epoch batch metrics.
• 🧭 SetFit Review Router: Utterance classification categories (Dietary/Service/General), Ollama GLM 5.1 augmentation, training, and evaluation.
• 🛠️ Troubleshooting Guide: Solutions for common issues like majority-class collapses, vocabulary matches, or scheduling.

Development

This project uses uv for dependency management.

# Install uv (if not already installed)
curl -LsSf https://astral.sh/uv/install.sh | sh

# Install dependencies (CPU-only PyTorch, no Ray)
uv sync

# Install with Ray distributed features
uv sync --extra ray

# Install dev and test suites
uv sync --extra dev --extra ray

# Install with diffusion (image generation) support
uv sync --extra diffusion

# Full development install
uv sync --extra dev --extra ray --extra diffusion

Local CUDA / GPU development

The locked packages resolve CPU-only PyTorch. To install CUDA-enabled PyTorch locally:

uv sync --no-sources-package torch

Note: This ignores CPU overrides in pyproject.toml and pulls PyTorch from PyPI with CUDA/NVIDIA libraries. Avoid committing changes to uv.lock.

Testing

Ensure local CI tests pass prior to submitting changes:

# Run all tests
uv run pytest tests/

# Run only Ray Train tests
uv run pytest tests/ -k "Ray"

# Run with coverage report
uv run pytest tests/ --cov=sentimentizer --cov-report=term-missing

# Verbose per-test output (or: make test-verbose)
uv run pytest tests/ -v

Project Structure

sentimentizer/
├── __init__.py          # Logging and timing utilities
├── compat.py            # Transformers/setfit compatibility shims
├── config.py            # Configuration dataclasses and constants
├── data_source.py       # Unified DataSource protocol (pandas/Ray)
├── device.py            # Device detection (cuda/mps/cpu)
├── env.py               # Environment setup (NVIDIA LD_LIBRARY_PATH)
├── extractor.py          # Ray Data extraction from zip/tar archives
├── exporter.py           # Standalone Prometheus metrics exporter
├── export_onnx.py        # ONNX export, quantization, validation
├── hf.py                # Hugging Face Hub push/pull + model card generation
├── hf_dataset.py        # Dataset wrapper and collation for HF transformers
├── hf_tokenizer.py      # Tokenizer wrapper for HF transformers
├── loader.py             # Data loading utilities
├── losses.py             # FocalCrossEntropyLoss for 3-class training
├── metrics.py            # 3-class classification metrics (per-class P/R/F1, balanced accuracy, MCC)
├── metrics_publisher.py   # Epoch metrics publishing (Prometheus + JSON) + intra-epoch batch snapshots
├── predictor.py           # SentimentPredictor (model loading, inference)
├── safety.py              # Shared prompt safety (NSFW blocklist, injection patterns)
├── serve/                 # Ray Serve deployment: FastAPI + @serve.ingress, /v1/ prefix
│   ├── app.py             # FastAPI route handlers and deployment class
│   ├── base.py            # ServiceMetrics (request/latency tracking), _DummyServe fallback
│   ├── config.py           # Serve deployment configuration (YAML/env var loading, incl. cors_origins)
│   ├── middleware.py       # Auth, rate limiting, idempotency, prompt safety for image routes
│   ├── models.py          # Pydantic request/response models for Swagger docs
│   ├── diffusion_models.py # Pydantic request/response models for image generation (+ Job models)
│   └── diffusion_app.py    # Headless ComfyUI deployment, dispatcher, and job endpoints
├── diffusion/                # ComfyUI workflows and image job support
│   ├── comfyui.py            # HTTP client + Krea 2 / Ideogram 4 native workflows
│   ├── config.py             # ImageModelConfig + YAML/env loading
│   ├── diffusion_config.yaml # ConvRot checkpoint names and sampling defaults
│   └── job_store.py          # JobStoreLogic + Ray actor for async job metadata
├── tokenizer.py           # Text tokenizer with pre-trained support
├── trainer.py             # Training logic
├── tuner.py               # Ray Tune + Optuna hyperparameter search
├── data/                  # Training data (Yelp, GloVe)
├── agent/                 # LLM-guided tuning agent
│   ├── __init__.py       # Package exports
│   ├── config.yaml       # Agent + tuner configuration (YAML)
│   ├── loader.py         # YAML → dataclass config loader
│   ├── models.py         # Pydantic models (AnalysisResult, TuningDecision, etc.)
│   ├── agents.py         # Pydantic AI agents (GLM 5.1 via Ollama)
│   ├── prompts.py        # System prompts for analysis & strategy agents
│   ├── state.py          # LangGraph AgentState TypedDict
│   ├── nodes.py          # LangGraph node functions (analyze, decide, tune, evaluate)
│   ├── graph.py          # LangGraph StateGraph + run_agent_tuning() entry point
│   └── diagnose_model.py # TuningRun workflow (tune → train → validate → retry pipeline)
├── router/                # SetFit router module
│   ├── __init__.py       # Package exports
│   ├── config.py         # SetFitConfig, RouteLabels, AugmentConfig
│   ├── seeds.py          # Golden example utterances per category
│   ├── augment.py        # GLM 5.1 augmentation via Ollama
│   ├── dataset.py        # JSONL dataset loader, train/test split
│   ├── train_router.py   # SetFit training with compat shims
│   └── evaluate.py       # Similarity heatmap, threshold calibration
└── models/
    ├── __init__.py
    ├── base.py            # BaseSentimentModel with predict() and predict_text()
    ├── hf_base.py         # Base class for Hugging Face transformer architectures
    ├── rnn.py            # Bidirectional LSTM (3-class output)
    ├── encoder.py         # Transformer encoder model (3-class output)
    ├── decoder.py         # Encoder-decoder transformer (3-class output)
    └── modernbert.py      # ModernBERT transformer classifier wrapper (3-class output)

License

MIT