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.21.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.21-py3-none-any.whl (25.8 kB view details)

Uploaded Python 3

File details

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

File metadata

File hashes

Hashes for water_conflict_classifier-0.1.21.tar.gz
Algorithm Hash digest
SHA256 72bb1c99dfe9139c67586d31c0c2ecae527547d48af115dfcca5c4cb88469786
MD5 2efcbb4d272ec9c45111845e79fe8540
BLAKE2b-256 06dadfeee00cf4106c164ef1c9754a989d8d4d76af7e0e0593be4e99ed2f1d16

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for water_conflict_classifier-0.1.21-py3-none-any.whl
Algorithm Hash digest
SHA256 6d83d1c6def505cb842b6c19745aeb8cc4467895fb62529decfb2f61604cab6c
MD5 d41c5ee454c0855cbb6294f851285875
BLAKE2b-256 d52b48ce91d5b5e557780d3b473de9257091f6f92f86fa9ca96689b8b73b9195

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