Skip to main content

SetFit-based multi-label classifier for water-related conflict events

Project description

Water Conflict Classifier

SetFit-based multi-label text classifier for identifying water-related conflict events in news headlines.

This folder contains the package source code. For usage instructions with the published package, see the PyPI page.

Project: Experimental research supporting the Pacific Institute's Water Conflict Chronology
Developer: Baobab Tech
License: CC BY-NC 4.0 (Non-Commercial)
PyPI Package: water-conflict-classifier


Package Installation & Usage

Install from PyPI:

pip install water-conflict-classifier

Use the trained model:

from setfit import SetFitModel

model = SetFitModel.from_pretrained("baobabtech/water-conflict-classifier")
predictions = model.predict(["Military group attack workers at dam"])
# Returns: [[1, 1, 1]]  # [Trigger, Casualty, Weapon]

The rest of this README is for developers who want to train their own model or modify the package.

Frugal AI: Training with Limited Data

This classifier demonstrates an intentional approach to building AI systems with limited data using SetFit - a framework for few-shot learning with sentence transformers. Rather than defaulting to massive language models (GPT, Claude, or 100B+ parameter models) for simple classification tasks, we fine-tune small, efficient models (e.g., BAAI/bge-small-en-v1.5 with ~33M parameters) on a focused dataset.

Why this matters: The industry has normalized using trillion-parameter models to classify headlines, answer simple questions, or categorize text - tasks that don't require world knowledge, reasoning, or generative capabilities. This is computationally wasteful and environmentally costly. A properly fine-tuned small model can achieve comparable or better accuracy while using a fraction of the compute resources.

Our approach:

  • Train on ~600 examples (few-shot learning with SetFit)
  • Deploy small parameter models (e.g., ~33M params) vs. 100B-1T parameter alternatives
  • Achieve specialized task performance without the overhead of general-purpose LLMs
  • Reduce inference costs and latency by orders of magnitude

This is not about avoiding large models altogether - they're invaluable for complex reasoning tasks. But for targeted classification problems with labeled data, fine-tuning remains the professional, responsible choice.

Package Structure

This is the source code for the water-conflict-classifier Python package, published to PyPI.

classifier/
├── __init__.py                         # Package marker
├── data_prep.py                        # Data loading (for training-ready datasets)
├── training_logic.py                   # Core training logic (fixed hyperparameters)
├── training_logic_optuna.py            # Optuna hyperparameter search training
├── evaluation.py                       # Model evaluation & metrics
├── model_card.py                       # Model card generation
├── versioning.py                       # Experiment tracking & versioning
├── evals_upload.py                     # Upload evaluation results to HF
├── train_setfit_headline_classifier.py # Local training script
├── pyproject.toml                      # Package configuration
├── setup.py                            # Build configuration
└── README.md                           # This file

Note: Scripts that use this package (like cloud training with HF Jobs and dataset preparation) are in the ../scripts/ folder.


Local Training

Train on your own hardware with local data files.

Setup

  1. Install the package in development mode:
uv pip install -e .
# or with regular pip:
pip install -e .

This installs the modules (data_prep, training_logic, etc.) so they can be imported.

  1. Prepare training data:

Training data should be in ../data/:

  • ../data/positives.csv - Water conflict headlines with labels
  • ../data/negatives.csv - Non-water conflict headlines

Generate negatives from ACLED (if needed):

cd ../scripts
python transform_prep_negatives.py
  1. Train:
python train_setfit_headline_classifier.py

Model saved to ./water-conflict-classifier/


Cloud Training

For training on HuggingFace Jobs (managed GPUs):

  1. Prepare training dataset: Use ../scripts/prepare_training_dataset.py to preprocess, balance, and upload training-ready data to HF Hub
  2. Train model: Two options:
    • ../scripts/train_on_hf.py - Standard training with fixed hyperparameters
    • ../scripts/train_on_hf_optuna.py - Optuna hyperparameter search (finds optimal settings automatically)

Optuna training searches on a small sample (~200 examples) for speed, then trains the final model on full data with best hyperparameters. Results are logged to the HF evals dataset for comparison.

See the scripts README for complete workflow.


Publishing to PyPI

See PUBLISHING.md for complete instructions on building and publishing the package.


Data Sources

The training data combines:

  • Positive Examples: Water conflict headlines from Pacific Institute Water Conflict Chronology
  • Negative Examples: Two types for balanced training:
    1. Hard Negatives (~120): Water-related peaceful news (infrastructure, research, conservation) to prevent false positives
    2. ACLED Negatives (~600): Non-water conflict events from ACLED

Hard Negatives Strategy

Without hard negatives, the model learns "water mentioned → conflict" instead of "water + violence → conflict". Hard negatives are water-related headlines that lack violence or conflict:

  • Water infrastructure projects (dams, treatment plants)
  • Scientific water research and technology
  • Water conservation initiatives and conferences
  • Environmental water management

These are tagged with priority_sample=True in the dataset and are ALWAYS included in training (never diluted by sampling). This ensures the model correctly distinguishes peaceful water news from actual water conflicts.

Resources


License

Copyright © 2025 Baobab Tech

This project is licensed under the Creative Commons Attribution-NonCommercial 4.0 International License.

You are free to use, share, and adapt this work for non-commercial purposes with appropriate attribution to Baobab Tech. For commercial licensing inquiries, please contact Baobab Tech.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

water_conflict_classifier-0.1.22.tar.gz (26.3 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

water_conflict_classifier-0.1.22-py3-none-any.whl (25.9 kB view details)

Uploaded Python 3

File details

Details for the file water_conflict_classifier-0.1.22.tar.gz.

File metadata

File hashes

Hashes for water_conflict_classifier-0.1.22.tar.gz
Algorithm Hash digest
SHA256 fc9f329ef8a663b11b36be8d47f46004543ab7900fcb61e9322187ce6705d41a
MD5 dcd73ffee597b23074f41228e3c182df
BLAKE2b-256 9c0f0622c1a567f26b2a74c0cb0251b8ac75988d4fa2d18701c5d536697875b9

See more details on using hashes here.

File details

Details for the file water_conflict_classifier-0.1.22-py3-none-any.whl.

File metadata

File hashes

Hashes for water_conflict_classifier-0.1.22-py3-none-any.whl
Algorithm Hash digest
SHA256 c526a595f4f05260eb13f0eff6f3b0260428b9b7d1213d167ffc0836ade188fd
MD5 aa7f119b52a03aa4b92d7b7bc9285492
BLAKE2b-256 1286f4aab96f432757e520e23076e25225b2ecedd07b7b65952ed7459e28d0b0

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page