Sync from GitHub e9a00339708d14fd15e64679f58d1c321b37ea9d

.github/workflows/ci.yml

@@ -0,0 +1,25 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+          cache: "pip"
+
+      - name: Install dependencies
+        run: pip install -r requirements.txt
+
+      - name: Run tests
+        run: pytest src/ -v

.github/workflows/sync-to-hf-space.yml

@@ -0,0 +1,24 @@
+name: Sync to Hugging Face Space
+
+on:
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+jobs:
+  sync:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Push to Hugging Face Space
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: |
+          git config user.name "github-actions[bot]"
+          git config user.email "github-actions[bot]@users.noreply.github.com"
+          git remote add space https://Evasim:${HF_TOKEN}@huggingface.co/spaces/Evasim/Simple-Text-Classifier
+          git checkout --orphan space-sync
+          git add -A
+          git commit -m "Sync from GitHub ${GITHUB_SHA}"
+          git push --force space space-sync:main

.gitignore

@@ -0,0 +1,233 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#   Usually these files are written by a python script from a template
+#   before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+# Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+# uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+# poetry.lock
+# poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+# pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Data (regenerate via src/data_ingestion.py)
+data/*.csv
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#   and can be added to the global gitignore or merged into this file.  For a more nuclear
+#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
+# .idea/
+
+# Abstra
+#   Abstra is an AI-powered process automation framework.
+#   Ignore directories containing user credentials, local state, and settings.
+#   Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#   and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#   you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+# Temporary file for partial code execution
+tempCodeRunnerFile.py
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml
+
+efs/
+*.pt
+*.pkl
+*.onnx
+
+venv/
+.venv/
+__pycache__
+
+.DS_Store
+.gradio/

.python-version

@@ -0,0 +1 @@
+3.10.11

README.md

@@ -0,0 +1,46 @@
+---
+title: Simple Text Classifier
+emoji: 📰
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 6.24.0
+app_file: app.py
+pinned: false
+license: mit
+short_description: Classifies news into World, Sports, Business, or Sci/Tech.
+---
+
+# My first MLOps Project
+
+[![CI](https://github.com/evasim/my-first-MLOPS-project/actions/workflows/ci.yml/badge.svg)](https://github.com/evasim/my-first-MLOPS-project/actions/workflows/ci.yml)
+
+🚧 Currently just started to learn MLOps (how to build and deploy models) — following the [Made With ML](https://madewithml.com/) course.
+
+📊 Using one of the public datasets on Hugging Face, [fancyzhx/ag_news](https://huggingface.co/datasets/fancyzhx/ag_news), to classify news headlines into World, Sports, Business, or Sci/Tech.
+
+🚀 **Live demo:** [huggingface.co/spaces/Evasim/Simple-Text-Classifier](https://huggingface.co/spaces/Evasim/Simple-Text-Classifier)
+
+🤗 **Trained model:** [huggingface.co/Evasim/First_Project](https://huggingface.co/Evasim/First_Project) (fine-tuned SciBERT)
+
+## Reproducing locally
+
+```bash
+git clone https://github.com/evasim/my-first-MLOPS-project.git
+cd my-first-MLOPS-project
+pip install -r requirements.txt
+
+# regenerate the dataset (not committed — it's public and regenerable)
+python src/data_ingestion.py
+
+# run the test suite
+pytest src/
+```
+
+To run predictions against the trained model, add to a `.env` file:
+```
+HF_REPO_ID=Evasim/First_Project
+```
+Then: `python -m src.predict "your headline here"`
+
+Training happens in [notebooks/MLOPS_project.ipynb](notebooks/MLOPS_project.ipynb).

app.py

@@ -0,0 +1,23 @@
+import gradio as gr
+from dotenv import load_dotenv
+
+from src.predict import predict
+
+load_dotenv()
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(lines=3, placeholder="Enter a news headline or article", label="Text"),
+    outputs=gr.Textbox(label="Predicted category"),
+    title="AG News Classifier",
+    description="Fine-tuned SciBERT model that classifies text into World, Sports, Business, or Sci/Tech.",
+    examples=[
+        "Wall Street stocks rallied after the Federal Reserve meeting.",
+        "The team won the championship game last night.",
+        "Scientists discover a new exoplanet orbiting a distant star.",
+        "World leaders meet to discuss climate change policy.",
+    ],
+)
+
+if __name__ == "__main__":
+    demo.launch(share=True)

notebooks/MLOPS_project.ipynb

@@ -0,0 +1,650 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install \"ray[default]\"\n",
+    "%pip install python-dotenv"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os \n",
+    "import ray "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": "import sys; sys.path.append(\"..\")\nimport warnings; warnings.filterwarnings(\"ignore\")\nfrom dotenv import load_dotenv; load_dotenv(override=True)\n%load_ext autoreload\n%autoreload 2"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": "if ray.is_initialized():\n    ray.shutdown()\nray.init(\n    num_cpus= 4,\n    object_store_memory=2 * 1024 * 1024 * 1024,\n    runtime_env={\"working_dir\": \"..\", \"env_vars\": {\"USE_LIBUV\": \"0\"}}\n)"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ray.cluster_resources()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "num_workers = 2\n",
+    "resources_per_worker = {\"CPU\": 1, \"GPU\": 1}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "if os.path.exists(\"/efs\"):\n",
+    "    EFS_DIR = f\"/efs/shared_storage/MY-FIRST-MLOPS-PROJECT/{os.environ.get('Eva', 'default_user')}\"\n",
+    "else:    \n",
+    "    EFS_DIR = str(Path(os.getcwd()).parent / \"local_storage\")\n",
+    "    os.makedirs(EFS_DIR, exist_ok=True)\n",
+    "\n",
+    "print(f\"Current active storage directory: {EFS_DIR}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": "from pathlib import Path\ndataset = str(Path.cwd().parent / \"data\" / \"raw_dataset.csv\")\ndf = pd.read_csv(dataset)\ndf.head()"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sklearn.model_selection import train_test_split"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df.label.value_counts()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_size = 0.2 \n",
+    "train_df, val_df = train_test_split(df, stratify=df.label, test_size=test_size, random_state=42)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_df.label.value_counts()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "val_df.label.value_counts() * int((1 - test_size) / test_size)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections import Counter \n",
+    "import matplotlib.pyplot as plt \n",
+    "import seaborn as sns; sns.set_theme()\n",
+    "import warnings; warnings.filterwarnings(\"ignore\")\n",
+    "from wordcloud import WordCloud, STOPWORDS"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tags = Counter(df.label)\n",
+    "tags.most_common()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "28",
+   "metadata": {},
+   "outputs": [],
+   "source": "from src.data import data, stratify_split, preprocess\nray.data.DataContext.get_current().execution_options.preserve_order = True"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "33",
+   "metadata": {},
+   "outputs": [],
+   "source": "import os \nimport random\nimport numpy as np\nimport torch \nfrom ray.data.preprocessor import Preprocessor\n\ndef setting_seeds(seed=42):\n    np.random.seed(seed)\n    random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    eval(\"setattr(torch.backends.cudnn, 'deterministic', True)\") # forces GPU to always perform calculation in exact sequence\n    eval(\"setattr(torch.backends.cudnn, 'benchmark', False)\") \n    os.environ[\"PYTHONHASHSEED\"] = str(seed)"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "35",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CustomPreprocessor():\n",
+    "    \"\"\"Custom Preprocess.\"\"\"\n",
+    "    def __init__(self, label_decoder = None):\n",
+    "        self.label_decoder = label_decoder or {\n",
+    "            0: \"World\",\n",
+    "            1: \"Sports\",\n",
+    "            2: \"Business\",\n",
+    "            3: \"Sci/Tech\"\n",
+    "        }\n",
+    "        self.class_to_index = {v:k for k, v in self.label_decoder.items()}\n",
+    "\n",
+    "    def fitting(self, ds):\n",
+    "        _ = ds.unique(column=\"label\")\n",
+    "        return self\n",
+    "    \n",
+    "    def transforming(self, ds):\n",
+    "        return ds.map_batches(\n",
+    "            preprocess, \n",
+    "            batch_format = \"pandas\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "36",
+   "metadata": {},
+   "outputs": [],
+   "source": "import json\nimport torch\nimport torch.nn as nn\nfrom transformers import BertModel\nfrom transformers import AutoTokenizer"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "37",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_name = \"allenai/scibert_scivocab_uncased\"\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+    "llm = BertModel.from_pretrained(model_name, return_dict = False)\n",
+    "embedding_dim = llm.config.hidden_size\n",
+    "num_classes = 4"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "38",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FinetunedLLM(nn.Module):\n",
+    "    def __init__(self, llm, dropout_p, embedding_dim, num_classes):\n",
+    "        super(FinetunedLLM, self).__init__()\n",
+    "        self.llm = llm\n",
+    "        self.dropout_p = dropout_p\n",
+    "        self.embedding_dim = embedding_dim\n",
+    "        self.num_classes = num_classes\n",
+    "        self.dropout = torch.nn.Dropout(dropout_p)\n",
+    "        self.fc1 =torch.nn.Linear(embedding_dim, num_classes)\n",
+    "\n",
+    "    def forward(self, batch):\n",
+    "        ids, masks = batch[\"ids\"], batch[\"masks\"]\n",
+    "        seq, pool = self.llm(input_ids = ids, attention_mask = masks)\n",
+    "        z = self.dropout(pool)\n",
+    "        z = self.fc1(z)\n",
+    "        return z\n",
+    "    \n",
+    "    @torch.inference_mode()\n",
+    "    def predict(self,batch):\n",
+    "        self.eval()\n",
+    "        z = self(batch)\n",
+    "        y_pred = torch.argmax(z, dim=1).cpu().numpy()\n",
+    "        return y_pred\n",
+    "    \n",
+    "    def save(self, dp):\n",
+    "        with open(Path(dp, \"args.json\"), \"w\") as fp:\n",
+    "            contents = {\n",
+    "                \"dropout_p\": self.dropout_p,\n",
+    "                \"embedding_dim\": self.embedding_dim,\n",
+    "                \"num_classes\": self.num_classes,\n",
+    "            }\n",
+    "            json.dump(contents, fp, indent=4, sort_keys=False)\n",
+    "        torch.save(self.state_dict(), os.path.join(dp, \"model.pt\"))\n",
+    "\n",
+    "    @classmethod\n",
+    "    def load(cls, args_fp, state_dict_fp):\n",
+    "        with open(args_fp, \"r\") as fp:\n",
+    "            kwargs = json.load(fp = fp)\n",
+    "        llm = BertModel.from_pretrained(model_name, return_dict = False)\n",
+    "        model = cls(llm = llm, **kwargs)\n",
+    "        model.load_state_dict(torch.load(state_dict_fp, map_location=torch.device(\"cpu\")))\n",
+    "        return model "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "39",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = FinetunedLLM(llm=llm, dropout_p=0.5, embedding_dim=embedding_dim, num_classes=num_classes)\n",
+    "print(model.named_parameters)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "40",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ray.train.torch import get_device"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "41",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def pad_array(arr, dtype=np.int32):\n",
+    "    max_len = max(len(row) for row in arr)\n",
+    "    padded_arr = np.zeros((arr.shape[0], max_len), dtype=dtype)\n",
+    "    for i, row in enumerate(arr):\n",
+    "        padded_arr[i][:len(row)] = row\n",
+    "    return padded_arr"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "42",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def collate_fn(batch):\n",
+    "    batch[\"ids\"] = pad_array(batch[\"ids\"])\n",
+    "    batch[\"masks\"] = pad_array(batch[\"masks\"])\n",
+    "    dtypes = {\"ids\": torch.int32, \"masks\": torch.int32, \"targets\": torch.int64}\n",
+    "    tensor_batch = {}\n",
+    "    for key, array in batch.items():\n",
+    "        tensor_batch[key] = torch.as_tensor(array, dtype=dtypes[key], device=get_device())\n",
+    "    return tensor_batch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "43",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path \n",
+    "import ray.train as train\n",
+    "from ray.train import Checkpoint, CheckpointConfig, DataConfig, RunConfig, ScalingConfig\n",
+    "from ray.train.torch import TorchCheckpoint, TorchTrainer, TorchConfig\n",
+    "import tempfile\n",
+    "import torch.nn.functional as F\n",
+    "from torch.nn.parallel.distributed import DistributedDataParallel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "44",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train_step(ds, batch_size, model, num_classes, loss_fn, optimizer):\n",
+    "    model.train()\n",
+    "    loss = 0.0\n",
+    "    ds_generator = ds.iter_torch_batches(batch_size=batch_size, collate_fn=collate_fn)\n",
+    "    for i, batch in enumerate(ds_generator):\n",
+    "        optimizer.zero_grad()\n",
+    "        z = model(batch)\n",
+    "        targets = F.one_hot(batch[\"targets\"], num_classes=num_classes).float()\n",
+    "        J = loss_fn(z, targets)\n",
+    "        J.backward()\n",
+    "        optimizer.step()\n",
+    "        loss += (J.detach().item() - loss) / (i+1)\n",
+    "    return loss\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "45",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def eval_step(ds, batch_size, model, num_classes, loss_fn):\n",
+    "    model.eval()\n",
+    "    loss = 0.0 \n",
+    "    y_trues, y_preds = [], []\n",
+    "    ds_generator = ds.iter_torch_batches(batch_size=batch_size, collate_fn=collate_fn)\n",
+    "    with torch.inference_mode():\n",
+    "        for i, batch in enumerate(ds_generator):\n",
+    "            z = model(batch)\n",
+    "            targets = F.one_hot(batch[\"targets\"], num_classes=num_classes).float()\n",
+    "            J = loss_fn(z, targets).item()\n",
+    "            loss += (J-loss)/(i+1)\n",
+    "            y_trues.extend(batch[\"targets\"].cpu().numpy())\n",
+    "            y_preds.extend(torch.argmax(z, dim=1).cpu().numpy())\n",
+    "        return loss, np.vstack(y_trues), np.vstack(y_preds)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "46",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train_loop_per_worker(config):\n",
+    "    dropout_p = config[\"dropout_p\"]\n",
+    "    lr = config[\"lr\"]\n",
+    "    lr_factor = config[\"lr_factor\"]\n",
+    "    lr_patience = config[\"lr_patience\"]\n",
+    "    num_epochs = config[\"num_epochs\"]\n",
+    "    batch_size = config[\"batch_size\"]\n",
+    "    num_classes = config[\"num_classes\"]\n",
+    "\n",
+    "    setting_seeds()\n",
+    "    train_ds = train.get_dataset_shard(\"train\")\n",
+    "    val_ds = train.get_dataset_shard(\"val\")\n",
+    "    \n",
+    "    llm = BertModel.from_pretrained(\"allenai/scibert_scivocab_uncased\", return_dict = False)\n",
+    "    model = FinetunedLLM(llm=llm, dropout_p=dropout_p, embedding_dim=llm.config.hidden_size, num_classes=num_classes)\n",
+    "    model = train.torch.prepare_model(model)\n",
+    "\n",
+    "    loss_fn = nn.BCEWithLogitsLoss()\n",
+    "    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
+    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode=\"min\", factor=lr_factor, patience=lr_patience)\n",
+    "\n",
+    "    num_workers = train.get_context().get_world_size()\n",
+    "    batch_size_per_worker = batch_size // num_workers\n",
+    "    for epoch in range(num_epochs):\n",
+    "        train_loss = train_step(train_ds, batch_size_per_worker, model, num_classes, loss_fn, optimizer)\n",
+    "        val_loss, _, _ = eval_step(val_ds, batch_size_per_worker, model, num_classes, loss_fn)\n",
+    "        scheduler.step(val_loss)\n",
+    "\n",
+    "        with tempfile.TemporaryDirectory() as dp:\n",
+    "            if isinstance(model, DistributedDataParallel):\n",
+    "                model.module.save(dp=dp)\n",
+    "            else:\n",
+    "                model.save(dp=dp)\n",
+    "            metrics = dict(epoch=epoch, lr=optimizer.param_groups[0][\"lr\"], train_loss=train_loss, val_loss=val_loss)\n",
+    "            checkpoint = Checkpoint.from_directory(dp)\n",
+    "            train.report(metrics, checkpoint=checkpoint)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from src.config import EFS_DIR, BASE_DIR, DATA_DIR, ARTIFACTS_DIR"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "48",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_loop_config = {\n",
+    "    \"dropout_p\": 0.5,\n",
+    "    \"lr\": 1e-4,\n",
+    "    \"lr_factor\": 0.8,\n",
+    "    \"lr_patience\": 3,\n",
+    "    \"num_epochs\": 10,\n",
+    "    \"batch_size\": 256,\n",
+    "    \"num_classes\": num_classes,\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "49",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "scaling_config = ScalingConfig(\n",
+    "    num_workers=num_workers,\n",
+    "    use_gpu=bool(resources_per_worker[\"GPU\"]),\n",
+    "    resources_per_worker=resources_per_worker\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "50",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_config = CheckpointConfig(num_to_keep=1, checkpoint_score_attribute=\"val_loss\", checkpoint_score_order=\"min\")\n",
+    "run_config = RunConfig(name=\"llm\", checkpoint_config=checkpoint_config, storage_path=EFS_DIR)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "51",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds = data(dataset_loc=dataset)\n",
+    "train_ds, val_ds = stratify_split(ds, stratify=\"label\", test_size=test_size)\n",
+    "\n",
+    "preprocessor = CustomPreprocessor()\n",
+    "preprocessor = preprocessor.fitting(train_ds)\n",
+    "# materialize now, while the full cluster is free, so trainer.fit() doesn't have to\n",
+    "# tokenize on the fly while competing with the training worker for CPU\n",
+    "train_ds = preprocessor.transforming(train_ds).materialize()\n",
+    "val_ds = preprocessor.transforming(val_ds).materialize()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "52",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = TorchTrainer(\n",
+    "    train_loop_per_worker=train_loop_per_worker,\n",
+    "    train_loop_config=train_loop_config,\n",
+    "    scaling_config=scaling_config,\n",
+    "    run_config=run_config,\n",
+    "    datasets={\"train\": train_ds, \"val\": val_ds},\n",
+    "    dataset_config=DataConfig(datasets_to_split=[\"train\"]),\n",
+    "    torch_config=TorchConfig(backend=\"gloo\"),\n",
+    ")\n",
+    "results = trainer.fit()\n",
+    "results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "best_checkpoint = results.best_checkpoints[0][0]\n",
+    "best_metrics = results.best_checkpoints[0][1]\n",
+    "print(f\"best checkpoint: {best_checkpoint.path}\")\n",
+    "print(f\"best metrics: {best_metrics}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "55",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from huggingface_hub import create_repo, upload_folder\n",
+    "\n",
+    "# points directly at the epoch-0 checkpoint already saved to disk\n",
+    "# (currently the best val_loss of any epoch trained so far)\n",
+    "checkpoint_dir = str(EFS_DIR / \"llm\" / \"checkpoint_2026-08-10_13-03-12.020315\")\n",
+    "\n",
+    "HF_REPO_ID = os.environ[\"HF_REPO_ID\"]  # e.g. \"yourname/agnews-scibert-classifier\", set in .env\n",
+    "\n",
+    "create_repo(HF_REPO_ID, exist_ok=True, repo_type=\"model\", private=True)\n",
+    "upload_folder(\n",
+    "    folder_path=checkpoint_dir,\n",
+    "    repo_id=HF_REPO_ID,\n",
+    "    repo_type=\"model\",\n",
+    "    commit_message=\"Upload epoch 0 checkpoint (best val_loss so far)\",\n",
+    ")\n",
+    "print(f\"Uploaded to https://huggingface.co/{HF_REPO_ID}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "57",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from huggingface_hub import snapshot_download\n",
+    "\n",
+    "HF_REPO_ID = os.environ[\"HF_REPO_ID\"]\n",
+    "eval_checkpoint_dir = snapshot_download(repo_id=HF_REPO_ID, repo_type=\"model\")\n",
+    "print(f\"Downloaded checkpoint to {eval_checkpoint_dir}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "58",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "eval_model = FinetunedLLM.load(\n",
+    "    args_fp=Path(eval_checkpoint_dir, \"args.json\"),\n",
+    "    state_dict_fp=Path(eval_checkpoint_dir, \"model.pt\"),\n",
+    ")\n",
+    "eval_device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "eval_model = eval_model.to(eval_device)\n",
+    "print(f\"Loaded checkpoint onto {eval_device}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "59",
+   "metadata": {},
+   "outputs": [],
+   "source": "def eval_collate_fn(batch):\n    batch[\"ids\"] = pad_array(batch[\"ids\"])\n    batch[\"masks\"] = pad_array(batch[\"masks\"])\n    dtypes = {\"ids\": torch.int32, \"masks\": torch.int32, \"targets\": torch.int64}\n    return {key: torch.as_tensor(array, dtype=dtypes[key], device=eval_device) for key, array in batch.items()}\n\ndef evaluate_checkpoint(ds, batch_size, model, num_classes):\n    model.eval()\n    loss_fn = nn.BCEWithLogitsLoss()\n    loss = 0.0\n    y_trues, y_preds = [], []\n    ds_generator = ds.iter_torch_batches(batch_size=batch_size, collate_fn=eval_collate_fn)\n    model_device = next(model.parameters()).device\n    with torch.inference_mode():\n        for i, batch in enumerate(ds_generator):\n            batch = {key: value.to(model_device) for key, value in batch.items()}\n            z = model(batch)\n            targets = F.one_hot(batch[\"targets\"], num_classes=num_classes).float()\n            J = loss_fn(z, targets).item()\n            loss += (J - loss) / (i + 1)\n            y_trues.extend(batch[\"targets\"].cpu().numpy())\n            y_preds.extend(torch.argmax(z, dim=1).cpu().numpy())\n    return loss, np.array(y_trues), np.array(y_preds)"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "60",
+   "metadata": {},
+   "outputs": [],
+   "source": "from sklearn.metrics import classification_report\n\nval_loss, y_true, y_pred = evaluate_checkpoint(val_ds, batch_size=64, model=eval_model, num_classes=num_classes)\nprint(f\"val_loss: {val_loss:.4f}\")\nprint(classification_report(y_true, y_pred, target_names=[preprocessor.label_decoder[i] for i in sorted(preprocessor.label_decoder)]))"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

requirements.txt

@@ -0,0 +1,18 @@
+ray[default]==2.56.1
+torch==2.11.0
+transformers==5.15.0
+datasets==5.0.1
+huggingface_hub==1.27.0
+python-dotenv==1.2.2
+pandas==2.3.3
+numpy==2.2.6
+scikit-learn==1.7.2
+nltk==3.10.2
+matplotlib==3.10.9
+seaborn==0.13.2
+wordcloud==1.9.6
+pyarrow==25.0.0
+tqdm==4.70.0
+pytest==8.3.4
+gradio==6.24.0
+spaces

src/config.py

@@ -0,0 +1,191 @@
+import os 
+from pathlib import Path
+
+SRC_DIR = Path(__file__).parent.absolute()
+BASE_DIR = SRC_DIR.parent
+DATA_DIR = BASE_DIR / "data"
+EFS_DIR = Path(os.getenv("EFS_DIR", BASE_DIR / "efs"))
+ARTIFACTS_DIR = EFS_DIR / "artifacts"
+ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
+
+STOPWORDS = [
+    "i",
+    "me",
+    "my",
+    "myself",
+    "we",
+    "our",
+    "ours",
+    "ourselves",
+    "you",
+    "you're",
+    "you've",
+    "you'll",
+    "you'd",
+    "your",
+    "yours",
+    "yourself",
+    "yourselves",
+    "he",
+    "him",
+    "his",
+    "himself",
+    "she",
+    "she's",
+    "her",
+    "hers",
+    "herself",
+    "it",
+    "it's",
+    "its",
+    "itself",
+    "they",
+    "them",
+    "their",
+    "theirs",
+    "themselves",
+    "what",
+    "which",
+    "who",
+    "whom",
+    "this",
+    "that",
+    "that'll",
+    "these",
+    "those",
+    "am",
+    "is",
+    "are",
+    "was",
+    "were",
+    "be",
+    "been",
+    "being",
+    "have",
+    "has",
+    "had",
+    "having",
+    "do",
+    "does",
+    "did",
+    "doing",
+    "a",
+    "an",
+    "the",
+    "and",
+    "but",
+    "if",
+    "or",
+    "because",
+    "as",
+    "until",
+    "while",
+    "of",
+    "at",
+    "by",
+    "for",
+    "with",
+    "about",
+    "against",
+    "between",
+    "into",
+    "through",
+    "during",
+    "before",
+    "after",
+    "above",
+    "below",
+    "to",
+    "from",
+    "up",
+    "down",
+    "in",
+    "out",
+    "on",
+    "off",
+    "over",
+    "under",
+    "again",
+    "further",
+    "then",
+    "once",
+    "here",
+    "there",
+    "when",
+    "where",
+    "why",
+    "how",
+    "all",
+    "any",
+    "both",
+    "each",
+    "few",
+    "more",
+    "most",
+    "other",
+    "some",
+    "such",
+    "no",
+    "nor",
+    "not",
+    "only",
+    "own",
+    "same",
+    "so",
+    "than",
+    "too",
+    "very",
+    "s",
+    "t",
+    "can",
+    "will",
+    "just",
+    "don",
+    "don't",
+    "should",
+    "should've",
+    "now",
+    "d",
+    "ll",
+    "m",
+    "o",
+    "re",
+    "ve",
+    "y",
+    "ain",
+    "aren",
+    "aren't",
+    "couldn",
+    "couldn't",
+    "didn",
+    "didn't",
+    "doesn",
+    "doesn't",
+    "hadn",
+    "hadn't",
+    "hasn",
+    "hasn't",
+    "haven",
+    "haven't",
+    "isn",
+    "isn't",
+    "ma",
+    "mightn",
+    "mightn't",
+    "mustn",
+    "mustn't",
+    "needn",
+    "needn't",
+    "shan",
+    "shan't",
+    "shouldn",
+    "shouldn't",
+    "wasn",
+    "wasn't",
+    "weren",
+    "weren't",
+    "won",
+    "won't",
+    "wouldn",
+    "wouldn't",
+]

src/data.py

@@ -0,0 +1,95 @@
+import re
+import ray
+import numpy as np
+import pandas as pd 
+from ray.data import Dataset
+from sklearn.model_selection import train_test_split
+from src.config import STOPWORDS 
+from typing import Dict,List,Tuple
+from transformers import BertTokenizer
+
+def data(dataset_loc: str, num_samples: int = None) -> ray.data.Dataset:
+    dataset = ray.data.read_csv(dataset_loc)
+    dataset = dataset.random_shuffle(seed = 1234)
+    dataset = ray.data.from_items(dataset.take(num_samples)) if num_samples else dataset
+    return dataset
+
+def stratify_split(
+        ds: Dataset, 
+        stratify: str, 
+        test_size: float, 
+        shuffle: bool = True,
+        seed: int = 1234,
+) -> Tuple[Dataset, Dataset]: 
+    def _add_split(df: pd.DataFrame) -> pd.DataFrame:
+        train, test = train_test_split(df, test_size=test_size, shuffle=shuffle, random_state=seed)
+        train["_split"]= "train"
+        test["_split"] = "test"
+        return pd.concat([train,test])
+    
+    def _filter_split(df: pd.DataFrame, split: str) -> pd.DataFrame:
+        return df[df["_split"] == split].drop("_split", axis = 1)
+    
+    grouped = ds.groupby(stratify).map_groups(_add_split, batch_format="pandas")
+    train_ds = grouped.map_batches(_filter_split, fn_kwargs={"split":"train"}, batch_format="pandas")
+    test_ds = grouped.map_batches(_filter_split, fn_kwargs={"split":"test"}, batch_format="pandas")
+
+    train_ds = train_ds.random_shuffle(seed=seed)
+    test_ds = test_ds.random_shuffle(seed=seed)
+
+    return train_ds, test_ds
+
+def clean_text(text: str) -> str:
+    # change every words into lower case
+    text = text.lower()
+
+    text = re.sub(r"http\S+", "", text) # remove links (must run before punctuation gets spaced apart below)
+
+    # removing stopwords such as "is", "the" and so on
+    pattern = re.compile(r'\b(' + r"|".join(STOPWORDS)+ r")\b\s*")
+    text = pattern.sub('', text)
+
+    text = re.sub(r"([!\"'#$%&()*\+,-./:;<=>?@\\\[\]^_`{|}~])", r" \1 ", text) # add space
+    text = re.sub("[^A-Za-z0-9]+", " ", text) # remove other than words and numbers
+    text = re.sub(" +", " ", text) # remove all extra spaces
+    text = text.strip() # remove spaces at the start and at the end
+
+    return text
+
+_tokenizer_cache = {}
+
+def _get_tokenizer(model_name: str = "allenai/scibert_scivocab_uncased"):
+    # cache per-process so map_batches workers don't reload the tokenizer on every batch
+    if model_name not in _tokenizer_cache:
+        _tokenizer_cache[model_name] = BertTokenizer.from_pretrained(model_name, return_dict=False)
+    return _tokenizer_cache[model_name]
+
+def tokenize(batch: Dict) -> Dict:
+    tokenizer = _get_tokenizer()
+    encoded = tokenizer(batch["text"].tolist(), return_tensors="np", padding = "longest")
+    return dict(ids=encoded["input_ids"], masks=encoded["attention_mask"], targets=np.array(batch["label"]))
+
+def preprocess(df: pd.DataFrame):
+    df["text"] = df.text.apply(clean_text)
+    targets = tokenize(df)
+    return targets 
+
+class CustomPreprocessor():
+    """Custom Preprocess."""
+    def __init__(self, label_decoder = None):
+        self.label_decoder = label_decoder or {
+            0: "World",
+            1: "Sports",
+            2: "Business",
+            3: "Sci/Tech"
+        }
+        self.class_to_index = {v:k for k, v in self.label_decoder.items()}
+
+    def fitting(self, ds):
+        _ = ds.unique(column="label")
+        return self
+    
+    def transforming(self, ds):
+        return ds.map_batches(
+            preprocess, 
+            batch_format = "pandas")

src/data_ingestion.py

@@ -0,0 +1,22 @@
+import os
+import pandas as pd
+from datasets import load_dataset
+
+def download_and_save_ag_news():
+    # Create data folder if missing
+    os.makedirs("data", exist_ok=True)
+    
+    print("⏳ Fetching ag_news from Hugging Face...")
+    # Load from HF Hub
+    dataset = load_dataset("fancyzhx/ag_news")
+    
+    # Convert training split to a DataFrame
+    train_df = pd.DataFrame(dataset["train"])
+    
+    # Save locally as a raw CSV
+    output_path = "data/raw_dataset.csv"
+    train_df.to_csv(output_path, index=False)
+    print(f"✅ Success! ag_news dataset saved locally to: {output_path}")
+
+if __name__ == "__main__":
+    download_and_save_ag_news()

src/models.py

@@ -0,0 +1,56 @@
+import json
+import os
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from transformers import BertModel
+
+MODEL_NAME = "allenai/scibert_scivocab_uncased"
+
+
+class FinetunedLLM(nn.Module):
+    def __init__(self, llm, dropout_p, embedding_dim, num_classes):
+        super().__init__()
+        self.llm = llm
+        self.dropout_p = dropout_p
+        self.embedding_dim = embedding_dim
+        self.num_classes = num_classes
+        self.dropout = nn.Dropout(dropout_p)
+        self.fc1 = nn.Linear(embedding_dim, num_classes)
+
+    def forward(self, batch):
+        ids, masks = batch["ids"], batch["masks"]
+        seq, pool = self.llm(input_ids=ids, attention_mask=masks)
+        z = self.dropout(pool)
+        z = self.fc1(z)
+        return z
+
+    @torch.inference_mode()
+    def predict(self, batch):
+        self.eval()
+        z = self(batch)
+        return torch.argmax(z, dim=1).cpu().numpy()
+
+    def save(self, dp):
+        with open(Path(dp, "args.json"), "w") as fp:
+            json.dump(
+                {
+                    "dropout_p": self.dropout_p,
+                    "embedding_dim": self.embedding_dim,
+                    "num_classes": self.num_classes,
+                },
+                fp,
+                indent=4,
+                sort_keys=False,
+            )
+        torch.save(self.state_dict(), os.path.join(dp, "model.pt"))
+
+    @classmethod
+    def load(cls, args_fp, state_dict_fp):
+        with open(args_fp, "r") as fp:
+            kwargs = json.load(fp)
+        llm = BertModel.from_pretrained(MODEL_NAME, return_dict=False)
+        model = cls(llm=llm, **kwargs)
+        model.load_state_dict(torch.load(state_dict_fp, map_location=torch.device("cpu")))
+        return model

src/predict.py

@@ -0,0 +1,49 @@
+import os
+import sys
+from pathlib import Path
+
+import spaces
+import torch
+from dotenv import load_dotenv
+from huggingface_hub import snapshot_download
+
+from src.data import _get_tokenizer, clean_text
+from src.models import FinetunedLLM
+
+LABEL_DECODER = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}
+
+_model_cache = {}
+
+
+def load_model(repo_id: str = None) -> FinetunedLLM:
+    """Download (if needed) and load the checkpoint, caching it per process."""
+    repo_id = repo_id or os.environ["HF_REPO_ID"]
+    if repo_id not in _model_cache:
+        checkpoint_dir = snapshot_download(repo_id=repo_id, repo_type="model")
+        model = FinetunedLLM.load(
+            args_fp=Path(checkpoint_dir, "args.json"),
+            state_dict_fp=Path(checkpoint_dir, "model.pt"),
+        )
+        model.eval()
+        _model_cache[repo_id] = model
+    return _model_cache[repo_id]
+
+
+@spaces.GPU
+def predict(text: str, repo_id: str = None) -> str:
+    model = load_model(repo_id)
+    tokenizer = _get_tokenizer()
+    encoded = tokenizer([clean_text(text)], return_tensors="pt", padding="longest")
+    batch = {"ids": encoded["input_ids"], "masks": encoded["attention_mask"]}
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+    batch = {k: v.to(device) for k, v in batch.items()}
+    with torch.inference_mode():
+        pred_idx = torch.argmax(model(batch), dim=1).item()
+    return LABEL_DECODER[pred_idx]
+
+
+if __name__ == "__main__":
+    load_dotenv()
+    text = " ".join(sys.argv[1:]) or "Wall Street stocks rallied after the Federal Reserve meeting."
+    print(predict(text))

src/test_data.py

@@ -0,0 +1,52 @@
+import numpy as np
+import pandas as pd
+import pytest
+import ray
+
+from src.data import clean_text, stratify_split, tokenize
+
+
+@pytest.mark.parametrize(
+    "text, expected",
+    [
+        ("Hello, World!", "hello world"),
+        ("Stocks   rally    after   the   Fed   meeting", "stocks rally fed meeting"),
+        ("Check https://example.com for details", "check details"),
+        ("", ""),
+    ],
+)
+def test_clean_text(text, expected):
+    assert clean_text(text) == expected
+
+
+def test_clean_text_is_lowercase():
+    assert clean_text("UPPER CASE TEXT") == clean_text("UPPER CASE TEXT").lower()
+
+
+def test_stratify_split_preserves_class_proportions():
+    df = pd.DataFrame(
+        {
+            "text": [f"sample {i}" for i in range(100)],
+            "label": [0] * 50 + [1] * 50,
+        }
+    )
+    ds = ray.data.from_pandas(df)
+    train_ds, test_ds = stratify_split(ds, stratify="label", test_size=0.2)
+
+    train_df = train_ds.to_pandas()
+    test_df = test_ds.to_pandas()
+
+    assert len(train_df) + len(test_df) == len(df)
+    assert set(train_df["label"].unique()) == {0, 1}
+    assert set(test_df["label"].unique()) == {0, 1}
+    assert train_df["label"].value_counts()[0] == train_df["label"].value_counts()[1]
+    assert test_df["label"].value_counts()[0] == test_df["label"].value_counts()[1]
+
+
+def test_tokenize_shapes_match_batch_size():
+    batch = pd.DataFrame({"text": ["hello world", "a longer sentence here"], "label": [0, 1]})
+    result = tokenize(batch)
+
+    assert result["ids"].shape[0] == len(batch)
+    assert result["masks"].shape == result["ids"].shape
+    assert np.array_equal(result["targets"], np.array([0, 1]))

src/test_predict.py

@@ -0,0 +1,14 @@
+import os
+
+import pytest
+
+from src.predict import LABEL_DECODER, predict
+
+pytestmark = pytest.mark.skipif(
+    "HF_REPO_ID" not in os.environ, reason="requires HF_REPO_ID (downloads the real checkpoint)"
+)
+
+
+def test_predict_returns_known_label():
+    label = predict("Wall Street stocks rallied after the Federal Reserve meeting.")
+    assert label in LABEL_DECODER.values()