Add CI, typed ORM models, and packaging cleanup

Introduce repo-wide engineering guardrails with GitHub Actions CI (ruff, mypy, bandit, pytest), pre-commit hooks, editor config, and contributor/security docs. Refactor API persistence models into `api/models` with SQLAlchemy 2.0 typed mappings, wire imports accordingly, and align project packaging/test execution by adding `py.typed` markers, pytest pythonpath config, and Makefile/docs updates that remove ad-hoc `PYTHONPATH` usage. This also cleans generated artifacts from git, updates notebook/script paths to src-layout module execution, and refreshes project structure documentation.

.editorconfig

@@ -0,0 +1,26 @@
+# EditorConfig — https://editorconfig.org
+root = true
+
+[*]
+charset = utf-8
+end_of_line = lf
+insert_final_newline = true
+trim_trailing_whitespace = true
+
+[*.{py,pyw}]
+indent_style = space
+indent_size = 4
+
+[*.{json,yml,yaml,toml}]
+indent_style = space
+indent_size = 2
+
+[Makefile]
+indent_style = tab
+
+[*.md]
+trim_trailing_whitespace = false
+
+[*.ipynb]
+indent_style = space
+indent_size = 2

.gitattributes

@@ -5,9 +5,7 @@
 # ==========================================
 # EXCLUSIONS FROM LANGUAGE STATISTICS
 # 1. Archived/Backup Dashboard Code
-# The dashboard_backup directory contains a legacy Vite/React application
-# including large minified JS assets in dist/, CSS, and vendored node_modules.
-# We exclude this entire directory to prevent archived code from dominating statistics.
+# If a legacy dashboard_backup directory is reintroduced, exclude generated assets.
 dashboard_backup/** linguist-generated=true
 # 2. IDE and Tooling Hooks
 # The .opencode directory contains generated JS hooks for the workspace environment.
@@ -17,7 +15,7 @@ dashboard_backup/** linguist-generated=true
 # Notebooks are used for exploration and training on Colab, but the core active
 # application is the Python backend and FastAPI. Excluding these prevents notebooks
 # from misrepresenting the primary languages.
-ml/notebooks/*.ipynb linguist-generated=true
+notebooks/*.ipynb linguist-generated=true
 # 4. Standard Build Artifacts & Vendored Dependencies
 # Ensures that any inadvertently committed build artifacts or vendor libraries
 # (like node_modules in subdirectories) do not skew language statistics.

.github/workflows/ci.yml

@@ -0,0 +1,78 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+concurrency:
+  group: ci-${{ github.ref }}
+  cancel-in-progress: true
+
+env:
+  PYTHON_VERSION: "3.11"
+
+jobs:
+  lint:
+    name: Lint (ruff)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ env.PYTHON_VERSION }}
+      - name: Install ruff
+        run: pip install ruff==0.15.11
+      - name: Ruff check
+        run: ruff check api src/absa tests scripts
+      - name: Ruff format check
+        run: ruff format --check api src/absa tests scripts
+
+  typecheck:
+    name: Typecheck (mypy)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ env.PYTHON_VERSION }}
+      - name: Install dev deps
+        run: |
+          pip install --upgrade pip
+          pip install -e ".[dev]"
+      - name: Mypy
+        run: mypy api src/absa
+
+  security:
+    name: Security (bandit)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ env.PYTHON_VERSION }}
+      - name: Install bandit
+        run: pip install bandit
+      - name: Bandit scan
+        run: bandit -r api src/absa
+
+  test:
+    name: Tests (pytest) — py${{ matrix.python }}
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python: ["3.10", "3.11"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+          cache: pip
+      - name: Install package + dev deps
+        run: |
+          pip install --upgrade pip
+          pip install -e ".[dev]"
+      - name: Run tests
+        run: pytest tests/ -v

.gitignore

@@ -14,6 +14,8 @@ __pycache__/
 .mypy_cache/
 .ruff_cache/
 .coverage
+.coverage.*
+coverage.xml
 htmlcov/
 
 # Virtual environment

.pre-commit-config.yaml

@@ -0,0 +1,33 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.11
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      - id: check-yaml
+        exclude: ^docker/docker-compose.*\.ya?ml$
+      - id: check-json
+      - id: end-of-file-fixer
+        exclude: ^notebooks/
+      - id: trailing-whitespace
+        exclude: ^notebooks/
+      - id: check-added-large-files
+        args: ["--maxkb=512"]  # ~data/model artifacts are DVC-tracked, not committed
+      - id: detect-private-key
+
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v3.17.0
+    hooks:
+      - id: pyupgrade
+        args: [--py310-plus]
+
+  - repo: https://github.com/econchick/interrogate
+    rev: 1.7.0
+    hooks:
+      - id: interrogate
+        args: [-vv, -i, --fail-under=20, --exclude-module=__init__, --ignore-module=__main__]

CONTRIBUTING.md

@@ -0,0 +1,89 @@
+# Contributing to Multilingual-ABSA
+
+Thanks for taking the time to contribute! This document outlines the workflow,
+tooling, and conventions for building and shipping changes to this repository.
+
+## Table of Contents
+
+- [Development Setup](#development-setup)
+- [Project Layout](#project-layout)
+- [Quality Gates](#quality-gates)
+- [Workflow](#workflow)
+- [Conventions](#conventions)
+- [Commit Guidelines](#commit-guidelines)
+
+## Development Setup
+
+```bash
+python -m venv .venv && source .venv/bin/activate
+pip install -e ".[dev]"
+```
+
+Install the pre-commit hooks (optional but recommended):
+
+```bash
+pre-commit install
+```
+
+## Project Layout
+
+```
+api/            # FastAPI REST service (routes, middleware, services, tasks, models, schemas)
+src/absa/       # Core ML library (data, models, evaluation, training, utils) — src-layout
+frontend/       # Streamlit dashboard
+scripts/        # Operational/one-off utility scripts
+notebooks/      # Exploration & Colab training notebooks
+tests/          # Pytest suite (api/, web/, unit/)
+docs/           # Project documentation
+docker/         # Container definitions & compose files
+monitoring/     # Prometheus / Grafana configuration
+data/           # Datasets (DVC-tracked)
+models/         # Model artifacts (DVC-tracked)
+```
+
+## Quality Gates
+
+Every change must pass all of the following before being merged:
+
+```bash
+make lint        # ruff check api src/absa tests
+make typecheck   # mypy api src/absa
+make security    # bandit -r api src/absa
+make test        # pytest
+```
+
+## Workflow
+
+1. **Fork** the repository and create a branch from `main`:
+
+   ```bash
+   git checkout -b feature/<description>
+   ```
+
+2. Make focused, atomic changes — see [Commit Guidelines](#commit-guidelines).
+
+3. Run the [quality gates](#quality-gates) locally.
+
+4. Open a pull request describing **what** changed, **why**, and how you
+   verified it. Reference any related issues.
+
+## Conventions
+
+- **Python** — target 3.10+. Format/lint is enforced by `ruff` (120-char lines,
+  `E`, `F`, `I`, `N`, `W` rule set). Type hints are checked by `mypy`.
+- **Imports** — absolute imports only (`from absa.data import ...`,
+  `from api.routes import ...`); never rely on `sys.path` hacks in library code.
+- **Models vs. schemas** — SQLAlchemy ORM models live in `api/models/`;
+  Pydantic request/response models live in `api/schemas/`.
+- **Secrets** — never commit `.env` or real credentials. Add any new required
+  environment variables to `.env.example`.
+- **Data** — datasets and model weights are versioned with DVC, not git.
+  Update `dvc.yaml` when preprocessing stages change.
+
+## Commit Guidelines
+
+- Keep commits small, focused, and logically independent.
+- Use the imperative mood: "Add batch status endpoint", not "Added endpoint".
+- Prefix with the area when it aids scanning, e.g. `api:`, `frontend:`,
+  `data:`, `docs:`.
+- Do not bundle unrelated changes (e.g. formatting + feature) in one commit.

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Multilingual-ABSA Contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Makefile

@@ -1,4 +1,4 @@
-.PHONY: install dev api frontend worker test lint typecheck security coverage docker-up docker-down clean
+.PHONY: install dev api frontend worker test lint format typecheck security coverage docker-up docker-down clean
 
 # ── Setup ────────────────────────────────────────────────────────────────
 install:
@@ -19,19 +19,22 @@ worker:
 
 # ── Quality gates ────────────────────────────────────────────────────────
 test:
-	PYTHONPATH=src pytest
+	pytest
 
 lint:
-	PYTHONPATH=src ruff check api src/absa tests
+	ruff check api src/absa tests
+
+format:
+	ruff format api src/absa tests
 
 typecheck:
-	PYTHONPATH=src mypy api src/absa
+	mypy api src/absa
 
 security:
 	bandit -r api src/absa
 
 coverage:
-	PYTHONPATH=src pytest --cov=api --cov=absa --cov-report=term-missing
+	pytest --cov=api --cov=absa --cov-report=term-missing
 
 # ── Docker ───────────────────────────────────────────────────────────────
 docker-up:

README.md

@@ -60,41 +60,42 @@ The system identifies **aspects** (specific features like "battery life", "sound
 
 ```
 .
-├── api/                        # FastAPI REST API (pure Python)
+├── api/                        # FastAPI REST service
 │   ├── main.py                 # Entry point — uvicorn api.main:app
-│   ├── middleware/             # Rate limiting, metrics, DB deps
-│   ├── routes/                 # /predict, /batch, /status, /health, /info
-│   ├── schemas/                # Pydantic models + SQLAlchemy ORM
+│   ├── models/                 # SQLAlchemy ORM models (Review, AspectResult, BatchJob)
+│   ├── middleware/             # Rate limiting, metrics (Prometheus), DB deps
+│   ├── routes/                 # /predict, /batch, /status, /download, /health, /info
+│   ├── schemas/                # Pydantic request/response models
 │   ├── services/               # ABSA inference pipeline, language detection
 │   └── tasks/                  # Celery batch processing workers
-├── src/absa/                   # Core ML library (pure Python)
+├── src/absa/                   # Core ML library (src-layout, pip-installable)
 │   ├── data/                   # Loading, preprocessing, augmentation, transliteration
 │   ├── models/                 # Training scripts (ONNX, Transformers, baselines)
 │   ├── evaluation/             # Cross-lingual eval, latency benchmarking
 │   ├── training/               # MLflow experiment tracking
-│   └── utils/                  # Path configuration
+│   └── utils/                  # Path + environment configuration
 ├── frontend/                   # Streamlit dashboard (pure Python, no HTML templates)
 │   ├── Home.py                 # Entry point — streamlit run frontend/Home.py
 │   ├── absa_client.py          # Thin HTTP client for the FastAPI backend
 │   ├── ui.py                   # Native Streamlit UI helpers
 │   └── views/                  # Pages: predict, admin (overview/batch/monitor)
-├── docker/                     # Containerisation
-│   ├── Dockerfile              # Production API image
-│   ├── Dockerfile.prod         # Production image (HuggingFace Hub model source)
-│   ├── docker-compose.yml      # Full stack: API + worker + DB + Redis + monitoring
-│   └── docker-compose.prod.yml # Production overrides
-├── tests/                      # Test suite
+├── tests/                      # Pytest suite
+│   ├── conftest.py             # sys.path bootstrap — no PYTHONPATH hacks required
 │   ├── api/                    # API endpoint tests
 │   ├── web/                    # Streamlit page + client tests
 │   └── unit/                   # Unit tests (bio tagger, lang detect)
-├── scripts/                    # Utility scripts
+├── scripts/                    # Operational utility scripts
+├── notebooks/                  # Exploration & Colab training notebooks (numbered 01–05)
 ├── monitoring/                 # Prometheus + Grafana config
 ├── docs/                       # Documentation
+├── docker/                     # Containerisation
 ├── data/                       # Datasets (managed by DVC)
-├── models/                     # ONNX model artifacts
+├── models/                     # ONNX model artifacts (DVC-tracked)
+├── .github/workflows/          # CI pipeline (lint, typecheck, security, tests)
 ├── .env.example                # Environment variable template
 ├── dvc.yaml                    # DVC data pipeline
-└── pyproject.toml              # Project metadata
+├── pyproject.toml              # Project metadata + tool config
+└── Makefile                    # Developer command shortcuts
 ```
 
 ---
@@ -177,7 +178,10 @@ dvc push       # Upload to remote storage
 ### Run Tests
 
 ```bash
-PYTHONPATH=.:src pytest tests/ -v
+make test        # pytest
+make lint        # ruff
+make typecheck   # mypy
+make security    # bandit
 ```
 
 ### Run Full Stack (Docker)

SECURITY.md

@@ -0,0 +1,31 @@
+# Security Policy
+
+## Reporting a Vulnerability
+
+Please **do not** open a public issue for security vulnerabilities. Instead,
+report them privately to the maintainers so they can be triaged and fixed
+before disclosure.
+
+Include in your report:
+
+- A description of the vulnerability and the affected endpoints/components.
+- Steps to reproduce (if possible).
+- Impact assessment.
+
+## Supported Versions
+
+Security fixes are backported to the latest stable release. Older versions are
+not actively patched — please upgrade to the current release.
+
+## Known Scope
+
+This project runs a public JSON API and an admin dashboard. A current threat
+model and mitigation checklist is maintained in
+[docs/SECURITY.md](docs/SECURITY.md) — please review it before deploying to an
+untrusted network.
+
+## Disclosure Timeline
+
+- **Acknowledgement** — within 72 hours of the report.
+- **Fix** — a patched release is published as soon as the fix is verified.
+- **Disclosure** — public mention of the vulnerability after the fix ships.

api/main.py

@@ -12,8 +12,8 @@ load_dotenv()
 
 from api.middleware.dependencies import engine  # noqa: E402
 from api.middleware.metrics import instrumentator  # noqa: E402
+from api.models.db_models import Base  # noqa: E402
 from api.routes import predict, results  # noqa: E402
-from api.schemas.db_models import Base  # noqa: E402
 from api.services.absa_pipeline import pipeline  # noqa: E402
 
 
@@ -41,7 +41,7 @@ app = FastAPI(
 )
 
 app.state.limiter = limiter
-app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)  # type: ignore[arg-type]
 
 app.add_middleware(
     CORSMiddleware,

api/middleware/dependencies.py

@@ -1,17 +1,15 @@
 import os
+
+from dotenv import load_dotenv
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
-from dotenv import load_dotenv
 
 load_dotenv()
 
 DATABASE_URL = os.getenv("DATABASE_URL")
 
 if not DATABASE_URL:
-    raise RuntimeError(
-        "DATABASE_URL environment variable is not set. "
-        "Please set it in your .env file or environment."
-    )
+    raise RuntimeError("DATABASE_URL environment variable is not set. Please set it in your .env file or environment.")
 
 connect_args = {}
 if DATABASE_URL.startswith("sqlite"):
@@ -21,7 +19,6 @@ engine = create_engine(DATABASE_URL, pool_pre_ping=True, connect_args=connect_ar
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
 
-
 def get_db():
     db = SessionLocal()
     try:

api/models/db_models.py

@@ -0,0 +1,55 @@
+"""SQLAlchemy ORM models for the API persistence layer.
+
+Typed with SQLAlchemy 2.0 ``Mapped`` annotations so mypy (via the
+``sqlalchemy.ext.mypy.plugin``) can infer attribute types instead of
+``Column[...]``.
+"""
+
+from __future__ import annotations
+
+import uuid
+from datetime import datetime, timezone
+
+from sqlalchemy import DateTime, Float, ForeignKey, Integer, String, Text
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+
+def _utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class Review(Base):
+    __tablename__ = "reviews"
+
+    id: Mapped[uuid.UUID] = mapped_column(primary_key=True, default=uuid.uuid4)
+    text: Mapped[str] = mapped_column(Text, nullable=False)
+    language: Mapped[str] = mapped_column(String(10), nullable=False)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow)
+    processing_time_ms: Mapped[float] = mapped_column(Float, nullable=False)
+
+
+class AspectResult(Base):
+    __tablename__ = "aspect_results"
+
+    id: Mapped[uuid.UUID] = mapped_column(primary_key=True, default=uuid.uuid4)
+    review_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("reviews.id"), nullable=False)
+    aspect: Mapped[str] = mapped_column(String(255), nullable=False)
+    sentiment: Mapped[str] = mapped_column(String(50), nullable=False)
+    confidence: Mapped[float] = mapped_column(Float, nullable=False)
+    start_pos: Mapped[int] = mapped_column(Integer, nullable=False)
+    end_pos: Mapped[int] = mapped_column(Integer, nullable=False)
+
+
+class BatchJob(Base):
+    __tablename__ = "batch_jobs"
+
+    id: Mapped[uuid.UUID] = mapped_column(primary_key=True, default=uuid.uuid4)
+    status: Mapped[str] = mapped_column(String(50), nullable=False, default="queued")
+    total: Mapped[int] = mapped_column(Integer, nullable=False)
+    processed: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow)
+    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

api/routes/predict.py

@@ -1,17 +1,19 @@
-from fastapi import APIRouter, Depends, HTTPException, UploadFile, File
-from sqlalchemy.orm import Session
-import pandas as pd
 import os
-import uuid
+import re
 import tempfile
 import time
-import re
+import uuid
+
+import pandas as pd
+from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
+from sqlalchemy.orm import Session
 
-from api.schemas.schemas import ReviewInput, PredictionResponse, BatchJobResponse
-from api.schemas.db_models import Review, AspectResult, BatchJob
 from api.middleware.dependencies import get_db
+from api.models.db_models import AspectResult, BatchJob, Review
+from api.schemas.schemas import BatchJobResponse, PredictionResponse, ReviewInput
 from api.services.absa_pipeline import pipeline
 from api.tasks.batch_tasks import process_batch
+
 router = APIRouter()
 
 
@@ -46,13 +48,13 @@ async def predict(request: ReviewInput, db: Session = Depends(get_db)):
         db.commit()
 
         return prediction
-    except Exception as e:
+    except Exception:
         raise HTTPException(status_code=500, detail="Model inference failed. Please try again.")
 
 
 @router.post("/batch", response_model=BatchJobResponse)
 async def predict_batch(file: UploadFile = File(...), db: Session = Depends(get_db)):
-    MAX_UPLOAD_SIZE = 50 * 1024 * 1024  # 50MB
+    max_upload_size = 50 * 1024 * 1024  # 50MB
 
     if not file.filename or not file.filename.endswith(".csv"):
         raise HTTPException(status_code=422, detail="Only CSV files are allowed.")
@@ -62,7 +64,7 @@ async def predict_batch(file: UploadFile = File(...), db: Session = Depends(get_
 
     try:
         content = await file.read()
-        if len(content) > MAX_UPLOAD_SIZE:
+        if len(content) > max_upload_size:
             raise HTTPException(status_code=422, detail="File exceeds 50MB maximum size.")
 
         # Create temp file to read
@@ -80,15 +82,11 @@ async def predict_batch(file: UploadFile = File(...), db: Session = Depends(get_
         df = pd.read_csv(tmp_path)
         if "text" not in df.columns:
             os.unlink(tmp_path)
-            raise HTTPException(
-                status_code=422, detail="CSV must contain a 'text' column."
-            )
+            raise HTTPException(status_code=422, detail="CSV must contain a 'text' column.")
 
         if len(df) > 10000:
             os.unlink(tmp_path)
-            raise HTTPException(
-                status_code=422, detail="Max 10,000 rows allowed per batch."
-            )
+            raise HTTPException(status_code=422, detail="Max 10,000 rows allowed per batch.")
 
         job_id_obj = uuid.uuid4()
         job_id = str(job_id_obj)
@@ -99,20 +97,16 @@ async def predict_batch(file: UploadFile = File(...), db: Session = Depends(get_
         # Queue Celery task
         process_batch.delay(job_id, tmp_path)
 
-        return BatchJobResponse(
-            job_id=job_id, status="queued", total_reviews=len(df), processed=0
-        )
+        return BatchJobResponse(job_id=job_id, status="queued", total_reviews=len(df), processed=0)
     except HTTPException:
         raise
     except Exception:
-        raise HTTPException(
-            status_code=500, detail="Batch processing failed. Please try again."
-        )
+        raise HTTPException(status_code=500, detail="Batch processing failed. Please try again.")
 
 
 @router.get("/status/{job_id}", response_model=BatchJobResponse)
 async def get_batch_status(job_id: str, db: Session = Depends(get_db)):
-    if not re.match(r'^[a-fA-F0-9\-]{36}$', job_id):
+    if not re.match(r"^[a-fA-F0-9\-]{36}$", job_id):
         raise HTTPException(status_code=400, detail="Invalid job ID format")
     try:
         job_id_uuid = uuid.UUID(job_id)

api/routes/results.py

@@ -1,9 +1,15 @@
-from fastapi import APIRouter
 import os
+import re
+from pathlib import Path
 from typing import Dict
 
+from fastapi import APIRouter, HTTPException
+from fastapi.responses import FileResponse
+
 router = APIRouter()
 
+_RESULTS_DIR = Path("data/results").resolve()
+
 
 @router.get("/health")
 async def health_check() -> Dict[str, str]:
@@ -20,25 +26,14 @@ async def get_info() -> Dict[str, str]:
         "max_batch_size": os.getenv("MAX_BATCH_SIZE", "10000"),
     }
 
-from fastapi import HTTPException
-from fastapi.responses import FileResponse
-from pathlib import Path
-import re
-import os
-
-_RESULTS_DIR = Path("data/results").resolve()
 
 @router.get("/download/{job_id}")
 async def download_result(job_id: str):
-    if not re.match(r'^[a-fA-F0-9\-]{36}$', job_id):
+    if not re.match(r"^[a-fA-F0-9\-]{36}$", job_id):
         raise HTTPException(status_code=400, detail="Invalid job ID format")
     resolved = (_RESULTS_DIR / f"{job_id}.csv").resolve()
     if not str(resolved).startswith(str(_RESULTS_DIR)):
         raise HTTPException(status_code=400, detail="Invalid job ID")
     if not resolved.exists():
         raise HTTPException(status_code=404, detail="Result file not found")
-    return FileResponse(
-        path=resolved,
-        filename=f"absa_results_{job_id}.csv",
-        media_type="text/csv"
-    )
+    return FileResponse(path=resolved, filename=f"absa_results_{job_id}.csv", media_type="text/csv")

api/schemas/db_models.py

@@ -1,43 +0,0 @@
-from sqlalchemy import Column, String, Integer, Float, DateTime, ForeignKey, Text, Uuid
-from sqlalchemy.orm import declarative_base
-import uuid
-from datetime import datetime, timezone
-
-Base = declarative_base()
-
-
-class Review(Base):
-    __tablename__ = "reviews"
-
-    id = Column(Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4)
-    text = Column(Text, nullable=False)
-    language = Column(String(10), nullable=False)
-    created_at = Column(
-        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
-    )
-    processing_time_ms = Column(Float, nullable=False)
-
-
-class AspectResult(Base):
-    __tablename__ = "aspect_results"
-
-    id = Column(Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4)
-    review_id = Column(Uuid(as_uuid=True), ForeignKey("reviews.id"), nullable=False)
-    aspect = Column(String(255), nullable=False)
-    sentiment = Column(String(50), nullable=False)
-    confidence = Column(Float, nullable=False)
-    start_pos = Column(Integer, nullable=False)
-    end_pos = Column(Integer, nullable=False)
-
-
-class BatchJob(Base):
-    __tablename__ = "batch_jobs"
-
-    id = Column(Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4)
-    status = Column(String(50), nullable=False, default="queued")
-    total = Column(Integer, nullable=False)
-    processed = Column(Integer, nullable=False, default=0)
-    created_at = Column(
-        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
-    )
-    completed_at = Column(DateTime(timezone=True), nullable=True)

api/schemas/schemas.py

@@ -1,5 +1,6 @@
+from typing import List, Optional
+
 from pydantic import BaseModel, ConfigDict, Field
-from typing import Optional, List
 
 
 class ReviewInput(BaseModel):

api/services/lang_service.py

@@ -1,6 +1,7 @@
-import fasttext
 from pathlib import Path
 
+import fasttext
+
 
 class LanguageService:
     def __init__(self):
@@ -15,7 +16,7 @@ class LanguageService:
     def detect_language(self, text: str) -> str:
         if self.model:
             predictions = self.model.predict(text.replace("\n", " "), k=1)
-            lang = predictions[0][0].replace("__label__", "")
+            lang: str = predictions[0][0].replace("__label__", "")
             if lang in ["en", "hi"]:
                 return lang
             # Default to en if unknown or other

api/tasks/__init__.py

@@ -1,11 +1,10 @@
-from celery import Celery
 import os
 
+from celery import Celery
+
 redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0")
 
-celery_app = Celery(
-    "absa_tasks", broker=redis_url, backend=redis_url.replace("/0", "/1")
-)
+celery_app = Celery("absa_tasks", broker=redis_url, backend=redis_url.replace("/0", "/1"))
 
 celery_app.conf.update(
     task_serializer="json",

api/tasks/batch_tasks.py

@@ -1,12 +1,14 @@
-from api.tasks import celery_app
-from api.services.absa_pipeline import pipeline
-from api.middleware.dependencies import SessionLocal
-from api.schemas.db_models import BatchJob, AspectResult, Review
-import pandas as pd
-import os
 import csv
+import os
 from datetime import datetime, timezone
 
+import pandas as pd
+
+from api.middleware.dependencies import SessionLocal
+from api.models.db_models import AspectResult, BatchJob, Review
+from api.services.absa_pipeline import pipeline
+from api.tasks import celery_app
+
 
 @celery_app.task(bind=True)
 def process_batch(self, job_id: str, file_path: str):
@@ -118,6 +120,7 @@ def process_batch(self, job_id: str, file_path: str):
             job.status = "failed"
             db.commit()
         import logging
+
         logging.exception("Batch processing failed for job %s", job_id)
     finally:
         # Clean up temp file

docs/DEPLOYMENT.md

@@ -27,8 +27,8 @@
 # Backend
 cp .env.example .env
 python -m venv .venv && source .venv/bin/activate
-pip install -r requirements.txt
-uvicorn app.main:app --reload --host 0.0.0.0 --port 8000
+pip install -e ".[dev]"
+uvicorn api.main:app --reload --host 0.0.0.0 --port 8000
 # API at http://localhost:8000, docs at http://localhost:8000/docs
 
 # MLflow

docs/TECH_STACK.md

@@ -18,25 +18,25 @@
 
 | Technology | Version | Purpose | Where Used |
 |------------|---------|---------|------------|
-| **PyTorch** | 2.3.0 | Deep learning framework | `src/models/` training |
+| **PyTorch** | 2.3.0 | Deep learning framework | `src/absa/models/` training |
 | **Transformers** | 4.39.3 | Model zoo, training, tokenization | All ML scripts |
 | **XLM-RoBERTa** | base | Multilingual encoder | `FacebookAI/xlm-roberta-base` |
 | **ONNX Runtime** | 1.18.0 | Production inference | `api/services/absa_pipeline.py` |
-| **Optimum** | 1.19.0 | ONNX export bridge | `src/models/export_onnx.py` |
+| **Optimum** | 1.19.0 | ONNX export bridge | `src/absa/models/export_onnx.py` |
 | **optimum-onnx** | (bundled) | ONNX runtime models | `ORTModelForTokenClassification`, `ORTModelForSequenceClassification` |
-| **PEFT** | 0.10.0 | Parameter-efficient fine-tuning | `src/models/train_qlora.py` (LoRA) |
-| **scikit-learn** | 1.4.2 | Metrics + baseline | `src/models/baseline.py`, `train_sentiment.py` |
-| **Datasets** | 2.19.0 | Data loading | `src/data/hf_dataset.py` |
-| **seqeval** | 1.2.2 | BIO tagging evaluation | `src/models/train_aspect_extraction.py` |
-| **fasttext-predict** | 0.9.2.4 | Language identification | `src/data/lang_detect.py`, `api/services/lang_service.py` |
-| **indic-nlp-library** | (git) | Devanagari transliteration | `src/data/transliterate.py` |
-| **nlpaug** | 1.1.11 | Text augmentation | `src/data/augmentation.py` |
+| **PEFT** | 0.10.0 | Parameter-efficient fine-tuning | `src/absa/models/train_qlora.py` (LoRA) |
+| **scikit-learn** | 1.4.2 | Metrics + baseline | `src/absa/models/baseline.py`, `train_sentiment.py` |
+| **Datasets** | 2.19.0 | Data loading | `src/absa/data/hf_dataset.py` |
+| **seqeval** | 1.2.2 | BIO tagging evaluation | `src/absa/models/train_aspect_extraction.py` |
+| **fasttext-predict** | 0.9.2.4 | Language identification | `src/absa/data/lang_detect.py`, `api/services/lang_service.py` |
+| **indic-nlp-library** | (git) | Devanagari transliteration | `src/absa/data/transliterate.py` |
+| **nlpaug** | 1.1.11 | Text augmentation | `src/absa/data/augmentation.py` |
 
 ## MLOps Stack
 
 | Technology | Version | Purpose | Where Used |
 |------------|---------|---------|------------|
-| **MLflow** | 2.13.0 | Experiment tracking | `src/training/mlflow_utils.py`, all `src/models/` |
+| **MLflow** | 2.13.0 | Experiment tracking | `src/absa/training/mlflow_utils.py`, all `src/absa/models/` |
 | **DVC** | 3.51.1 | Data version control | `dvc.yaml`, `.dvc/` |
 | **Evidently AI** | 0.4.30 | Data drift monitoring | `scripts/drift_monitor.py` |
 | **Prometheus** | latest | Metrics collection | `monitoring/prometheus.yml` |

notebooks/{03_train_colab.ipynb → 02_train_colab.ipynb}

@@ -17,7 +17,7 @@
       "source": [
         "!git clone https://github.com/Aryanmishra-dev/Multilingual-Absa.git\n",
         "%cd Multilingual-Absa\n",
-        "!pip install -r requirements.txt\n"
+        "!pip install .\n"
       ]
     },
     {
@@ -49,7 +49,7 @@
       "outputs": [],
       "source": [
         "# Prepare dataset\n",
-        "!PYTHONPATH=. python src/data/hf_dataset.py\n"
+        "!PYTHONPATH=src python -m absa.data.hf_dataset\n"
       ]
     },
     {
@@ -59,7 +59,7 @@
       "outputs": [],
       "source": [
         "# Run Aspect Extraction Training\n",
-        "!PYTHONPATH=. python src/models/train_aspect_extraction.py\n"
+        "!PYTHONPATH=src python -m absa.models.train_aspect_extraction\n"
       ]
     },
     {
@@ -69,7 +69,7 @@
       "outputs": [],
       "source": [
         "# Run Sentiment Classification Training\n",
-        "!PYTHONPATH=. python src/models/train_sentiment.py\n"
+        "!PYTHONPATH=src python -m absa.models.train_sentiment\n"
       ]
     },
     {
@@ -79,7 +79,7 @@
       "outputs": [],
       "source": [
         "# Run Baseline as well\n",
-        "!PYTHONPATH=. python src/models/baseline.py\n"
+        "!PYTHONPATH=src python -m absa.models.baseline\n"
       ]
     },
     {
@@ -89,7 +89,7 @@
       "outputs": [],
       "source": [
         "# Cross-lingual Evaluation\n",
-        "!PYTHONPATH=. python src/evaluation/cross_lingual_eval.py\n"
+        "!PYTHONPATH=src python -m absa.evaluation.cross_lingual_eval\n"
       ]
     },
     {

notebooks/README.md

@@ -0,0 +1,16 @@
+# Notebooks
+
+Exploratory analysis and Colab training notebooks. Run them in order. The
+first is best run locally; the training ones expect a Colab (T4+/A100) or
+CUDA GPU environment.
+
+| Notebook | Purpose |
+|----------|---------|
+| `01_data_exploration.ipynb`     | Dataset overview, language distribution, label stats |
+| `02_train_colab.ipynb`          | End-to-end training (aspect extraction + sentiment) on Colab |
+| `03_model_comparison.ipynb`     | Compare trained runs from the MLflow tracking server |
+| `04_qlora_colab.ipynb`          | QLoRA 4-bit fine-tuning of the sequence classifier |
+| `05_final_evaluation.ipynb`     | Final cross-lingual evaluation and latency benchmarks |
+
+These notebooks are generated/updated from `scripts/generate_notebooks.py`
+where applicable.

pyproject.toml

@@ -75,9 +75,14 @@ dev = [
 where = ["src", "."]
 include = ["absa*", "api*"]
 
+[tool.setuptools.package-data]
+absa = ["py.typed"]
+api = ["py.typed"]
+
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 python_files = ["test_*.py"]
+pythonpath = [".", "src"]
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
 addopts = "--cov=api --cov=absa --cov-report=term-missing --cov-report=xml"
@@ -92,6 +97,7 @@ extend-ignore = ["N999"]  # module file name style
 
 [tool.mypy]
 python_version = "3.10"
+plugins = ["sqlalchemy.ext.mypy.plugin"]
 warn_return_any = true
 warn_unused_configs = true
 ignore_missing_imports = true

scripts/README.md

@@ -0,0 +1,13 @@
+# Scripts
+
+Operational and one-off utility scripts. Most are invoked manually or via
+`Makefile`/`dvc.yaml` targets. Run from the repository root.
+
+| Script | Purpose |
+|--------|---------|
+| `download_data.py`         | Fetch fastText LID model, SemEval, and Amazon Hindi datasets |
+| `upload_models.py`         | Push trained model artifacts to HuggingFace Hub |
+| `generate_notebooks.py`    | Regenerate Colab notebook scaffolding |
+| `mlflow_ui.sh`             | Launch the MLflow tracking UI |
+| `drift_monitor.py`         | Evidently data-drift monitoring on live predictions |
+| `init_db.py`               | Create database tables from the SQLAlchemy models |

scripts/download_data.py

@@ -1,7 +1,10 @@
 import os
 import urllib.request
+
 from datasets import load_dataset
-from absa.utils.config import DATA_DIR, RAW_DIR, FASTTEXT_MODEL_PATH
+
+from absa.utils.config import FASTTEXT_MODEL_PATH, RAW_DIR
+
 
 def download_fasttext():
     url = "https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.ftz"
@@ -13,31 +16,33 @@ def download_fasttext():
         print("fastText LID model already exists.")
     print(f"fastText model size: {os.path.getsize(FASTTEXT_MODEL_PATH) / 1024 / 1024:.2f} MB")
 
+
 def download_semeval():
     print("Downloading SemEval datasets...")
     restaurants = load_dataset("tomaarsen/setfit-absa-semeval-restaurants")
     laptops = load_dataset("tomaarsen/setfit-absa-semeval-laptops")
-    
+
     rest_path = RAW_DIR / "semeval_restaurants"
     lap_path = RAW_DIR / "semeval_laptops"
-    
+
     restaurants.save_to_disk(str(rest_path))
     laptops.save_to_disk(str(lap_path))
-    
+
     print(f"SemEval Restaurants train samples: {len(restaurants['train'])}")
     print(f"SemEval Laptops train samples: {len(laptops['train'])}")
 
+
 def download_amazon_hindi():
     print("Downloading Amazon Hindi dataset...")
-    ds = load_dataset("ai4bharat/IndicSentiment", "translation-hi", 
-                      trust_remote_code=True, split="test")
-    
+    ds = load_dataset("ai4bharat/IndicSentiment", "translation-hi", trust_remote_code=True, split="test")
+
     amz_path = RAW_DIR / "amazon_hindi"
     os.makedirs(amz_path, exist_ok=True)
     file_path = amz_path / "hindi_sentiment.jsonl"
     ds.to_json(str(file_path))
     print(f"Downloaded {len(ds)} Hindi samples")
 
+
 if __name__ == "__main__":
     os.makedirs(RAW_DIR, exist_ok=True)
     download_fasttext()

scripts/drift_monitor.py

@@ -1,16 +1,17 @@
 import os
-import pandas as pd
 from datetime import datetime, timedelta
+
 import mlflow
-from evidently.report import Report
+import pandas as pd
 from evidently.metric_preset import DataDriftPreset, TextOverviewPreset
+from evidently.report import Report
 from sqlalchemy import create_engine
-import uuid
+
 
 def main():
     # Attempt to fetch database URL, fallback to sqlite for local tests
     db_url = os.getenv("DATABASE_URL", "sqlite:///./test.db")
-    
+
     # We would normally load the reference data (e.g. from training data CSV)
     # For this script, we'll assume a local path or create a dummy reference if missing
     ref_path = "data/reference.csv"
@@ -18,50 +19,49 @@ def main():
         ref_df = pd.read_csv(ref_path)
     else:
         print(f"Reference data not found at {ref_path}. Creating dummy reference data for testing.")
-        ref_df = pd.DataFrame({
-            "text": ["This is great", "I hate this", "Neutral statement"],
-            "language": ["en", "en", "en"]
-        })
-        
+        ref_df = pd.DataFrame(
+            {"text": ["This is great", "I hate this", "Neutral statement"], "language": ["en", "en", "en"]}
+        )
+
     try:
         # Load production data from the last 7 days
         engine = create_engine(db_url)
         seven_days_ago = datetime.now() - timedelta(days=7)
-        
+
         # Load directly from SQLAlchemy using pandas with parameterized query
         query = "SELECT text, language FROM reviews WHERE created_at >= %(cutoff)s"
         curr_df = pd.read_sql(query, engine, params={"cutoff": seven_days_ago})
     except Exception as e:
         print(f"Failed to fetch production data: {e}")
         curr_df = pd.DataFrame(columns=["text", "language"])
-        
+
     if len(curr_df) < 50:
-        print(f"Not enough data to run drift monitor (found {len(curr_df)} records, need at least 50). Exiting gracefully.")
+        print(
+            f"Not enough data to run drift monitor "
+            f"(found {len(curr_df)} records, need at least 50). Exiting gracefully."
+        )
         return
 
     # Run Evidently report
     print("Running Evidently drift report...")
-    report = Report(metrics=[
-        DataDriftPreset(),
-        TextOverviewPreset(column_name="text")
-    ])
-    
+    report = Report(metrics=[DataDriftPreset(), TextOverviewPreset(column_name="text")])
+
     report.run(reference_data=ref_df, current_data=curr_df)
-    
+
     # Create monitoring/reports dir if missing
     os.makedirs("monitoring/reports", exist_ok=True)
-    
+
     report_path = f"monitoring/reports/drift_{datetime.now().strftime('%Y%m%d')}.html"
     report.save_html(report_path)
     print(f"Report saved to {report_path}")
-    
+
     # Extract drift metrics as a dict
     report_dict = report.as_dict()
-    
+
     # Simplified check for drift (using Dataset Drift metric from DataDriftPreset)
     dataset_drift = report_dict["metrics"][0]["result"]["dataset_drift"]
     drift_share = report_dict["metrics"][0]["result"]["drift_share"]
-    
+
     if dataset_drift and drift_share > 0.3:
         print(f"⚠️ Drift detected — consider retraining. Drift share: {drift_share:.2f}")
         try:
@@ -72,5 +72,6 @@ def main():
         except Exception as e:
             print(f"Failed to log warning to MLflow: {e}")
 
+
 if __name__ == "__main__":
     main()

scripts/generate_notebooks.py

@@ -1,17 +1,20 @@
 import json
 from pathlib import Path
 
+
 def create_notebook(filename: str, cells_content: list):
     cells = []
     for content, cell_type in cells_content:
-        cells.append({
-            "cell_type": cell_type,
-            "metadata": {},
-            "execution_count": None if cell_type == "code" else None,
-            "outputs": [] if cell_type == "code" else None,
-            "source": [line + "\n" for line in content.split("\n")]
-        })
-        
+        cells.append(
+            {
+                "cell_type": cell_type,
+                "metadata": {},
+                "execution_count": None if cell_type == "code" else None,
+                "outputs": [] if cell_type == "code" else None,
+                "source": [line + "\n" for line in content.split("\n")],
+            }
+        )
+
         # Clean up outputs/execution_count for markdown
         if cell_type == "markdown":
             del cells[-1]["execution_count"]
@@ -20,11 +23,7 @@ def create_notebook(filename: str, cells_content: list):
     notebook = {
         "cells": cells,
         "metadata": {
-            "kernelspec": {
-                "display_name": "Python 3",
-                "language": "python",
-                "name": "python3"
-            },
+            "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
             "language_info": {
                 "codemirror_mode": {"name": "ipython", "version": 3},
                 "file_extension": ".py",
@@ -32,42 +31,116 @@ def create_notebook(filename: str, cells_content: list):
                 "name": "python",
                 "nbconvert_exporter": "python",
                 "pygments_lexer": "ipython3",
-                "version": "3.11.0"
-            }
+                "version": "3.11.0",
+            },
         },
         "nbformat": 4,
-        "nbformat_minor": 4
+        "nbformat_minor": 4,
     }
-    
+
     Path("notebooks").mkdir(parents=True, exist_ok=True)
     with open(f"notebooks/{filename}", "w") as f:
         json.dump(notebook, f, indent=2)
 
+
 def main():
     colab_cells = [
-        ("# Google Colab Training Notebook\n\nThis notebook is intended to be run on Google Colab with a T4 GPU. It clones the repo, installs dependencies, and runs the training scripts.", "markdown"),
-        ("!git clone https://github.com/Aryanmishra-dev/Multilingual-Absa.git\n%cd Multilingual-Absa\n!pip install .", "code"),
-        ("# Mount Google Drive to save models and MLflow logs persistently\nfrom google.colab import drive\ndrive.mount('/content/drive')", "code"),
-        ("# Create symlinks or copy data if needed\n# Assuming data is in the repo for now\n!mkdir -p /content/drive/MyDrive/ABSA_models", "code"),
+        (
+            "# Google Colab Training Notebook\n\n"
+            "This notebook is intended to be run on Google Colab with a T4 GPU. "
+            "It clones the repo, installs dependencies, and runs the training scripts.",
+            "markdown",
+        ),
+        (
+            "!git clone https://github.com/Aryanmishra-dev/Multilingual-Absa.git\n"
+            "%cd Multilingual-Absa\n!pip install .",
+            "code",
+        ),
+        (
+            "# Mount Google Drive to save models and MLflow logs persistently\n"
+            "from google.colab import drive\ndrive.mount('/content/drive')",
+            "code",
+        ),
+        (
+            "# Create symlinks or copy data if needed\n"
+            "# Assuming data is in the repo for now\n"
+            "!mkdir -p /content/drive/MyDrive/ABSA_models",
+            "code",
+        ),
         ("# Prepare dataset\n!PYTHONPATH=src python -m absa.data.hf_dataset", "code"),
-        ("# Run Aspect Extraction Training\n!PYTHONPATH=src python -m absa.models.train_aspect_extraction", "code"),
-        ("# Run Sentiment Classification Training\n!PYTHONPATH=src python -m absa.models.train_sentiment", "code"),
+        (
+            "# Run Aspect Extraction Training\n!PYTHONPATH=src python -m absa.models.train_aspect_extraction",
+            "code",
+        ),
+        (
+            "# Run Sentiment Classification Training\n!PYTHONPATH=src python -m absa.models.train_sentiment",
+            "code",
+        ),
         ("# Run Baseline as well\n!PYTHONPATH=src python -m absa.models.baseline", "code"),
-        ("# Cross-lingual Evaluation\n!PYTHONPATH=src python -m absa.evaluation.cross_lingual_eval", "code"),
-        ("# Copy models back to Drive\n!cp -r models/* /content/drive/MyDrive/ABSA_models/\n!cp -r mlflow /content/drive/MyDrive/ABSA_models/", "code")
+        (
+            "# Cross-lingual Evaluation\n!PYTHONPATH=src python -m absa.evaluation.cross_lingual_eval",
+            "code",
+        ),
+        (
+            "# Copy models back to Drive\n"
+            "!cp -r models/* /content/drive/MyDrive/ABSA_models/\n"
+            "!cp -r mlflow /content/drive/MyDrive/ABSA_models/",
+            "code",
+        ),
     ]
-    
+
     comparison_cells = [
-        ("# Model Comparison\n\nThis notebook connects to the MLflow tracking server and compares the results of our models.", "markdown"),
-        ("import mlflow\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport json\n\nmlflow.set_tracking_uri('sqlite:///mlflow/mlflow.db')", "code"),
-        ("# Load all runs\nexperiment = mlflow.get_experiment_by_name('multilingual-absa')\ndf = mlflow.search_runs(experiment_ids=[experiment.experiment_id])\ndisplay(df.head())", "code"),
-        ("# Bar chart: macro-F1 comparison\nmetrics = df[['tags.mlflow.runName', 'metrics.eval_macro_f1', 'metrics.test_f1', 'metrics.test_macro_f1', 'metrics.hindi_zero_shot_macro_f1']].fillna(0)\nmetrics['Best F1'] = metrics[['metrics.eval_macro_f1', 'metrics.test_f1', 'metrics.test_macro_f1']].max(axis=1)\n\nplt.figure(figsize=(10, 6))\nsns.barplot(data=metrics, x='tags.mlflow.runName', y='Best F1')\nplt.title('Model Comparison by Macro-F1 / Span-F1')\nplt.xticks(rotation=45)\nplt.show()", "code"),
-        ("# Load confusion matrix for best sentiment classifier\n# Note: Assuming the confusion_matrix.json artifact was downloaded or parsed.\nprint('Confusion Matrix (Placeholder for artifact loading)')", "code"),
-        ("# 5 Example Predictions\nprint('Example 1: The food was great but service was slow.')\nprint('Example 2: El sistema operativo es muy estable.')\nprint('... (Load pipeline and infer here)')", "code")
+        (
+            "# Model Comparison\n\n"
+            "This notebook connects to the MLflow tracking server and compares the "
+            "results of our models.",
+            "markdown",
+        ),
+        (
+            "import mlflow\nimport pandas as pd\nimport matplotlib.pyplot as plt\n"
+            "import seaborn as sns\nimport json\n\n"
+            "mlflow.set_tracking_uri('sqlite:///mlflow/mlflow.db')",
+            "code",
+        ),
+        (
+            "# Load all runs\n"
+            "experiment = mlflow.get_experiment_by_name('multilingual-absa')\n"
+            "df = mlflow.search_runs(experiment_ids=[experiment.experiment_id])\n"
+            "display(df.head())",
+            "code",
+        ),
+        (
+            "# Bar chart: macro-F1 comparison\n"
+            "metrics = df[['tags.mlflow.runName', 'metrics.eval_macro_f1', "
+            "'metrics.test_f1', 'metrics.test_macro_f1', "
+            "'metrics.hindi_zero_shot_macro_f1']].fillna(0)\n"
+            "metrics['Best F1'] = metrics[['metrics.eval_macro_f1', "
+            "'metrics.test_f1', 'metrics.test_macro_f1']].max(axis=1)\n\n"
+            "plt.figure(figsize=(10, 6))\n"
+            "sns.barplot(data=metrics, x='tags.mlflow.runName', y='Best F1')\n"
+            "plt.title('Model Comparison by Macro-F1 / Span-F1')\n"
+            "plt.xticks(rotation=45)\nplt.show()",
+            "code",
+        ),
+        (
+            "# Load confusion matrix for best sentiment classifier\n"
+            "# Note: Assuming the confusion_matrix.json artifact was downloaded "
+            "or parsed.\n"
+            "print('Confusion Matrix (Placeholder for artifact loading)')",
+            "code",
+        ),
+        (
+            "# 5 Example Predictions\n"
+            "print('Example 1: The food was great but service was slow.')\n"
+            "print('Example 2: El sistema operativo es muy estable.')\n"
+            "print('... (Load pipeline and infer here)')",
+            "code",
+        ),
     ]
-    
-    create_notebook("03_train_colab.ipynb", colab_cells)
+
+    create_notebook("02_train_colab.ipynb", colab_cells)
     create_notebook("03_model_comparison.ipynb", comparison_cells)
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     main()

scripts/init_db.py

@@ -1,11 +1,13 @@
 import os
-from sqlalchemy import create_engine
+import sys
+
 from dotenv import load_dotenv
+from sqlalchemy import create_engine
 
-import sys
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-from api.schemas.db_models import Base
+from api.models.db_models import Base
+
 
 def init_db():
     load_dotenv()
@@ -13,10 +15,11 @@ def init_db():
     if not database_url:
         print("DATABASE_URL not set in .env")
         return
-    
+
     engine = create_engine(database_url)
     Base.metadata.create_all(bind=engine)
     print("Database tables created successfully.")
 
+
 if __name__ == "__main__":
     init_db()

scripts/upload_models.py

@@ -1,6 +1,8 @@
 import os
+
 from huggingface_hub import HfApi
 
+
 def main():
     api = HfApi()
     repo_name = "multilingual-absa"
@@ -16,11 +18,10 @@ def main():
 
     print("Uploading models/onnx/ folder...")
     api.upload_folder(
-        folder_path="models/onnx/",
-        repo_id=repo_id,
-        commit_message="Upload INT8 ONNX models for Multilingual ABSA"
+        folder_path="models/onnx/", repo_id=repo_id, commit_message="Upload INT8 ONNX models for Multilingual ABSA"
     )
     print("Upload complete!")
 
+
 if __name__ == "__main__":
     main()

src/absa/data/augmentation.py

@@ -5,11 +5,12 @@ Targets minority classes in Hindi data (negative and conflict).
 
 import json
 import random
-from pathlib import Path
 from collections import Counter
-from transformers import MarianMTModel, MarianTokenizer
-import torch
+from pathlib import Path
+
 import mlflow
+import torch
+from transformers import MarianMTModel, MarianTokenizer
 
 random.seed(42)
 
@@ -33,12 +34,8 @@ class BackTranslator:
         return [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
 
     def back_translate(self, text):
-        en_translation = self.translate([text], self.hi2en_model, self.hi2en_tokenizer)[
-            0
-        ]
-        back_to_hi = self.translate(
-            [en_translation], self.en2hi_model, self.en2hi_tokenizer
-        )[0]
+        en_translation = self.translate([text], self.hi2en_model, self.hi2en_tokenizer)[0]
+        back_to_hi = self.translate([en_translation], self.en2hi_model, self.en2hi_tokenizer)[0]
         return back_to_hi
 
 

src/absa/data/bio_tagger.py

@@ -1,5 +1,5 @@
 import re
-from typing import List, Dict, Any, Tuple
+from typing import Any, Dict, List, Tuple
 
 
 def tokenize(text: str) -> List[Tuple[str, int, int]]:
@@ -20,9 +20,7 @@ def tokenize(text: str) -> List[Tuple[str, int, int]]:
     return tokens
 
 
-def convert_to_bio(
-    text: str, aspect_terms: List[Dict[str, Any]]
-) -> List[Dict[str, Any]]:
+def convert_to_bio(text: str, aspect_terms: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """
     Converts text and aspect spans to BIO tagged tokens.
 
@@ -58,11 +56,7 @@ def convert_to_bio(
             if not (t_end <= a_start or t_start >= a_end):
                 # There is overlap
                 # If this token overlaps with the start of the aspect
-                if t_start <= a_start or (
-                    len(bio_tags) > 0
-                    and bio_tags[-1]["label"] == "O"
-                    and t_start > a_start
-                ):
+                if t_start <= a_start or (len(bio_tags) > 0 and bio_tags[-1]["label"] == "O" and t_start > a_start):
                     label = "B-ASP"
                 else:
                     # Check if previous tag was B-ASP or I-ASP for the *same* aspect

src/absa/data/dataset.py

@@ -1,9 +1,11 @@
 import json
-from datasets import load_from_disk
 from collections import defaultdict
-from absa.utils.config import RAW_DIR, SEMEVAL_TRAIN_PATH, SEMEVAL_TEST_PATH
-from absa.data.preprocess import clean
+
+from datasets import load_from_disk
+
 from absa.data.lang_detect import detect_language
+from absa.data.preprocess import clean
+from absa.utils.config import RAW_DIR, SEMEVAL_TEST_PATH, SEMEVAL_TRAIN_PATH
 
 
 def process_semeval():
@@ -13,8 +15,8 @@ def process_semeval():
     rest_data = load_from_disk(str(rest_path))
     lap_data = load_from_disk(str(lap_path))
 
-    train_samples = defaultdict(list)
-    test_samples = defaultdict(list)
+    train_samples: dict[tuple[str, str], list[dict[str, str]]] = defaultdict(list)
+    test_samples: dict[tuple[str, str], list[dict[str, str]]] = defaultdict(list)
 
     for ds_name, ds, source_name in [
         ("train", rest_data["train"], "restaurants"),
@@ -37,7 +39,7 @@ def process_semeval():
         (SEMEVAL_TEST_PATH, test_samples),
     ]:
         total = 0
-        lang_counts = defaultdict(int)
+        lang_counts: defaultdict[str, int] = defaultdict(int)
         with open(path, "w", encoding="utf-8") as f:
             for (text, source), aspects in data_dict.items():
                 lang = detect_language(text)

src/absa/data/hf_dataset.py

@@ -1,11 +1,13 @@
 import json
+from pathlib import Path
+from typing import Any, Dict, List
+
 import numpy as np
 import pandas as pd
-from pathlib import Path
-from typing import List, Dict, Any
 from datasets import Dataset, DatasetDict
-from transformers import AutoTokenizer
 from sklearn.model_selection import train_test_split
+from transformers import AutoTokenizer
+
 from absa.data.bio_tagger import convert_to_bio
 
 np.random.seed(42)
@@ -147,12 +149,8 @@ def main():
     cls_df = pd.DataFrame(cls_data)
 
     # Stratified split 80/10/10 based on label
-    train_cls, temp_cls = train_test_split(
-        cls_df, test_size=0.2, random_state=42, stratify=cls_df["label"]
-    )
-    val_cls, test_cls = train_test_split(
-        temp_cls, test_size=0.5, random_state=42, stratify=temp_cls["label"]
-    )
+    train_cls, temp_cls = train_test_split(cls_df, test_size=0.2, random_state=42, stratify=cls_df["label"])
+    val_cls, test_cls = train_test_split(temp_cls, test_size=0.5, random_state=42, stratify=temp_cls["label"])
 
     def tokenize_cls(examples):
         # Format: [CLS] text [SEP] aspect_term [SEP]
@@ -172,9 +170,7 @@ def main():
         }
     )
 
-    tokenized_cls = cls_dataset.map(
-        tokenize_cls, batched=True, remove_columns=["text", "aspect_term", "id"]
-    )
+    tokenized_cls = cls_dataset.map(tokenize_cls, batched=True, remove_columns=["text", "aspect_term", "id"])
     tokenized_cls.save_to_disk(str(output_dir / "absa_cls_dataset"))
     print(f"CLS Dataset saved to {output_dir / 'absa_cls_dataset'}")
 

src/absa/data/hindi_loader.py

@@ -1,7 +1,8 @@
 import json
-from absa.utils.config import RAW_DIR, AMAZON_HINDI_PATH
-from absa.data.preprocess import clean
+
 from absa.data.lang_detect import detect_language
+from absa.data.preprocess import clean
+from absa.utils.config import AMAZON_HINDI_PATH, RAW_DIR
 
 
 def process_hindi():
@@ -15,9 +16,7 @@ def process_hindi():
     total = 0
     lang_counts = {"hi": 0, "hinglish": 0, "en": 0, "other": 0}
 
-    with open(raw_path, "r", encoding="utf-8") as fin, open(
-        AMAZON_HINDI_PATH, "w", encoding="utf-8"
-    ) as fout:
+    with open(raw_path, "r", encoding="utf-8") as fin, open(AMAZON_HINDI_PATH, "w", encoding="utf-8") as fout:
         for line in fin:
             row = json.loads(line)
             text = row.get("INDIC REVIEW", row.get("text", ""))

src/absa/data/lang_detect.py

@@ -1,5 +1,7 @@
 import re
+
 import fasttext
+
 from absa.utils.config import FASTTEXT_MODEL_PATH
 
 _model = None

src/absa/data/preprocess.py

@@ -1,5 +1,6 @@
 import re
 import unicodedata
+
 from absa.data.transliterate import transliterate
 
 

src/absa/data/transliterate.py

@@ -1,6 +1,6 @@
 import logging
-import unicodedata
 import re
+import unicodedata
 
 logger = logging.getLogger(__name__)
 
@@ -10,9 +10,7 @@ try:
     HAS_INDIC_NLP = True
 except ImportError:
     HAS_INDIC_NLP = False
-    logger.warning(
-        "indic-nlp-library not found. Transliteration will fallback to basic unicode handling."
-    )
+    logger.warning("indic-nlp-library not found. Transliteration will fallback to basic unicode handling.")
 
 
 def transliterate(text: str, src_lang: str) -> str:
@@ -30,11 +28,7 @@ def transliterate(text: str, src_lang: str) -> str:
             roman_text = text
     else:
         # Fallback to basic unicode normalization
-        roman_text = (
-            unicodedata.normalize("NFKD", text)
-            .encode("ascii", "ignore")
-            .decode("utf-8")
-        )
+        roman_text = unicodedata.normalize("NFKD", text).encode("ascii", "ignore").decode("utf-8")
 
     if not roman_text:
         roman_text = text

src/absa/evaluation/benchmark_latency.py

@@ -4,9 +4,10 @@ Script to benchmark latency for PyTorch, ONNX, and ONNX INT8 models on CPU.
 
 import time
 from pathlib import Path
+
 import numpy as np
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 try:
     from optimum.onnxruntime import ORTModelForSequenceClassification
@@ -30,9 +31,7 @@ def benchmark_model(model, tokenizer, texts, model_type="pytorch"):
 
     print(f"Benchmarking {model_type}...")
     for text in texts:
-        inputs = tokenizer(
-            [text], return_tensors="pt", padding=True, truncation=True, max_length=128
-        )
+        inputs = tokenizer([text], return_tensors="pt", padding=True, truncation=True, max_length=128)
 
         start_time = time.perf_counter()
         if model_type == "pytorch":
@@ -82,12 +81,8 @@ def main():
         # 2. ONNX CPU
         if onnx_dir.exists():
             print("Loading ONNX model...")
-            onnx_model = ORTModelForSequenceClassification.from_pretrained(
-                str(onnx_dir)
-            )
-            mean_onnx, p95_onnx, tput_onnx = benchmark_model(
-                onnx_model, tokenizer, texts, "onnx"
-            )
+            onnx_model = ORTModelForSequenceClassification.from_pretrained(str(onnx_dir))
+            mean_onnx, p95_onnx, tput_onnx = benchmark_model(onnx_model, tokenizer, texts, "onnx")
             results["ONNX (CPU)"] = {
                 "mean_ms": mean_onnx,
                 "p95_ms": p95_onnx,
@@ -97,12 +92,8 @@ def main():
         # 3. ONNX INT8 CPU
         if int8_dir.exists():
             print("Loading ONNX INT8 model...")
-            int8_model = ORTModelForSequenceClassification.from_pretrained(
-                str(int8_dir)
-            )
-            mean_int8, p95_int8, tput_int8 = benchmark_model(
-                int8_model, tokenizer, texts, "onnx_int8"
-            )
+            int8_model = ORTModelForSequenceClassification.from_pretrained(str(int8_dir))
+            mean_int8, p95_int8, tput_int8 = benchmark_model(int8_model, tokenizer, texts, "onnx_int8")
             results["ONNX INT8 (CPU)"] = {
                 "mean_ms": mean_int8,
                 "p95_ms": p95_int8,
@@ -110,26 +101,18 @@ def main():
             }
 
     print("\n--- Latency Benchmark Results ---")
-    print(
-        f"{'Model':<20} | {'Mean (ms)':<10} | {'P95 (ms)':<10} | {'Throughput (samples/s)':<25}"
-    )
+    print(f"{'Model':<20} | {'Mean (ms)':<10} | {'P95 (ms)':<10} | {'Throughput (samples/s)':<25}")
     print("-" * 75)
     for name, metrics in results.items():
-        print(
-            f"{name:<20} | {metrics['mean_ms']:<10.2f} | {metrics['p95_ms']:<10.2f} | {metrics['throughput']:<25.2f}"
-        )
+        print(f"{name:<20} | {metrics['mean_ms']:<10.2f} | {metrics['p95_ms']:<10.2f} | {metrics['throughput']:<25.2f}")
 
     # Target check
     if "ONNX INT8 (CPU)" in results:
         int8_p95 = results["ONNX INT8 (CPU)"]["p95_ms"]
         if int8_p95 < 300:
-            print(
-                f"\nSUCCESS: ONNX INT8 P95 latency is {int8_p95:.2f}ms, which is < 300ms target."
-            )
+            print(f"\nSUCCESS: ONNX INT8 P95 latency is {int8_p95:.2f}ms, which is < 300ms target.")
         else:
-            print(
-                f"\nWARNING: ONNX INT8 P95 latency is {int8_p95:.2f}ms, which is > 300ms target."
-            )
+            print(f"\nWARNING: ONNX INT8 P95 latency is {int8_p95:.2f}ms, which is > 300ms target.")
 
     mlflow.set_tracking_uri("sqlite:///mlflow.db")
     mlflow.set_experiment("latency-benchmark")

src/absa/evaluation/cross_lingual_eval.py

@@ -1,14 +1,15 @@
 import json
-import torch
 from pathlib import Path
+
+import mlflow
+import torch
+from sklearn.metrics import f1_score
 from transformers import (
-    AutoTokenizer,
-    AutoModelForTokenClassification,
     AutoModelForSequenceClassification,
+    AutoModelForTokenClassification,
+    AutoTokenizer,
     pipeline,
 )
-from sklearn.metrics import f1_score
-import mlflow
 
 from absa.training.mlflow_utils import setup_mlflow
 
@@ -30,20 +31,14 @@ def main():
     sentiment_model_path = Path("models/sentiment/best")
 
     if not aspect_model_path.exists() or not sentiment_model_path.exists():
-        print(
-            "Models not found locally. Skipping cross-lingual evaluation until models are trained."
-        )
+        print("Models not found locally. Skipping cross-lingual evaluation until models are trained.")
         return
 
     print("Loading models...")
     tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
 
-    aspect_model = AutoModelForTokenClassification.from_pretrained(
-        str(aspect_model_path)
-    )
-    sentiment_model = AutoModelForSequenceClassification.from_pretrained(
-        str(sentiment_model_path)
-    )
+    aspect_model = AutoModelForTokenClassification.from_pretrained(str(aspect_model_path))
+    sentiment_model = AutoModelForSequenceClassification.from_pretrained(str(sentiment_model_path))
 
     device = 0 if torch.cuda.is_available() else -1
 
@@ -79,9 +74,7 @@ def main():
             true_labels.append(sentiment_map[true_polarity])
 
             # Inference Sentiment
-            inputs = tokenizer(
-                text, term, return_tensors="pt", truncation=True, max_length=128
-            )
+            inputs = tokenizer(text, term, return_tensors="pt", truncation=True, max_length=128)
             if device == 0:
                 inputs = {k: v.to("cuda") for k, v in inputs.items()}
                 sentiment_model.to("cuda")
@@ -92,11 +85,7 @@ def main():
 
             pred_labels.append(pred_idx)
 
-    hindi_macro_f1 = (
-        f1_score(true_labels, pred_labels, average="macro")
-        if len(true_labels) > 0
-        else 0.0
-    )
+    hindi_macro_f1 = f1_score(true_labels, pred_labels, average="macro") if len(true_labels) > 0 else 0.0
     print(f"Hindi Zero-Shot Macro-F1: {hindi_macro_f1}")
 
     # We retrieve the best English test score from MLflow

src/absa/evaluation/final_eval.py

@@ -1,5 +1,5 @@
-import os
 import json
+import os
 
 
 def run_evaluation():

src/absa/models/baseline.py

@@ -1,13 +1,15 @@
 import json
-import joblib
 from pathlib import Path
 from typing import List
+
+import joblib
+import mlflow
 import pandas as pd
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.linear_model import LogisticRegression
-from sklearn.metrics import f1_score, confusion_matrix, classification_report
+from sklearn.metrics import classification_report, confusion_matrix, f1_score
 from sklearn.model_selection import train_test_split
-import mlflow
+
 from absa.training.mlflow_utils import log_training_run
 
 
@@ -60,47 +62,37 @@ def main():
     cls_df = extract_sentence_sentiment(train_df_raw)
 
     # Exact same split logic as hf_dataset.py
-    train_cls, temp_cls = train_test_split(
-        cls_df, test_size=0.2, random_state=42, stratify=cls_df["label"]
-    )
-    val_cls, test_cls = train_test_split(
-        temp_cls, test_size=0.5, random_state=42, stratify=temp_cls["label"]
-    )
-
-    X_train = train_cls["text"].values
+    train_cls, temp_cls = train_test_split(cls_df, test_size=0.2, random_state=42, stratify=cls_df["label"])
+    val_cls, test_cls = train_test_split(temp_cls, test_size=0.5, random_state=42, stratify=temp_cls["label"])
+
+    x_train = train_cls["text"].values
     y_train = train_cls["label"].values
 
-    X_test = test_cls["text"].values
+    x_test = test_cls["text"].values
     y_test = test_cls["label"].values
 
-    print(f"Training on {len(X_train)} samples, testing on {len(X_test)} samples.")
+    print(f"Training on {len(x_train)} samples, testing on {len(x_test)} samples.")
 
     # Baseline Model Pipeline
     vectorizer = TfidfVectorizer(ngram_range=(1, 2), max_features=10000)
-    classifier = LogisticRegression(
-        max_iter=1000, class_weight="balanced", random_state=42
-    )
+    classifier = LogisticRegression(max_iter=1000, class_weight="balanced", random_state=42)
 
     # Train
     print("Training TF-IDF + Logistic Regression...")
-    X_train_vec = vectorizer.fit_transform(X_train)
-    classifier.fit(X_train_vec, y_train)
+    x_train_vec = vectorizer.fit_transform(x_train)
+    classifier.fit(x_train_vec, y_train)
 
     # Evaluate
     print("Evaluating...")
-    X_test_vec = vectorizer.transform(X_test)
-    y_pred = classifier.predict(X_test_vec)
+    x_test_vec = vectorizer.transform(x_test)
+    y_pred = classifier.predict(x_test_vec)
 
     # Metrics
     macro_f1 = f1_score(y_test, y_pred, average="macro")
     per_class_f1 = f1_score(y_test, y_pred, average=None)
     conf_matrix = confusion_matrix(y_test, y_pred)
 
-    print(
-        classification_report(
-            y_test, y_pred, target_names=["positive", "negative", "neutral", "conflict"]
-        )
-    )
+    print(classification_report(y_test, y_pred, target_names=["positive", "negative", "neutral", "conflict"]))
 
     # Format metrics for MLflow
     metrics = {

src/absa/models/export_onnx.py

@@ -7,8 +7,8 @@ from pathlib import Path
 
 try:
     from optimum.onnxruntime import (
-        ORTModelForTokenClassification,
         ORTModelForSequenceClassification,
+        ORTModelForTokenClassification,
         ORTQuantizer,
     )
     from optimum.onnxruntime.configuration import AutoQuantizationConfig
@@ -19,9 +19,7 @@ except ImportError:
     print("Warning: optimum library not installed. Models will not be exported.")
 
 
-def export_and_quantize(
-    model_type: str, source_dir: Path, export_dir: Path, quantize_dir: Path
-):
+def export_and_quantize(model_type: str, source_dir: Path, export_dir: Path, quantize_dir: Path):
     print(f"Exporting {model_type} model from {source_dir} to {export_dir}")
 
     if not source_dir.exists():
@@ -35,13 +33,9 @@ def export_and_quantize(
     # when `export=True` is passed for HF models, it sets dynamic sequence lengths automatically.
 
     if model_type == "token_classification":
-        model = ORTModelForTokenClassification.from_pretrained(
-            str(source_dir), export=True
-        )
+        model = ORTModelForTokenClassification.from_pretrained(str(source_dir), export=True)
     elif model_type == "sequence_classification":
-        model = ORTModelForSequenceClassification.from_pretrained(
-            str(source_dir), export=True
-        )
+        model = ORTModelForSequenceClassification.from_pretrained(str(source_dir), export=True)
     else:
         raise ValueError(f"Unknown model_type: {model_type}")
 

src/absa/models/train_aspect_extraction.py

@@ -1,16 +1,17 @@
-import numpy as np
 from pathlib import Path
+
+import mlflow
+import numpy as np
 from datasets import load_from_disk
+from seqeval.metrics import f1_score as seqeval_f1_score
 from transformers import (
     AutoModelForTokenClassification,
-    TrainingArguments,
-    Trainer,
-    DataCollatorForTokenClassification,
     AutoTokenizer,
+    DataCollatorForTokenClassification,
+    Trainer,
+    TrainingArguments,
     set_seed,
 )
-from seqeval.metrics import f1_score as seqeval_f1_score
-import mlflow
 
 from absa.training.mlflow_utils import setup_mlflow
 
@@ -112,16 +113,13 @@ def main():
 
     # Log test metric manually since trainer.train() only automatically logs eval metrics
     # if report_to="mlflow" handles it, but test results we need to make sure are in the same run.
+    active_run = mlflow.active_run()
+    fallback_run_id = active_run.info.run_id if active_run else None
+
     with mlflow.start_run(
-        run_id=(
-            trainer.state.trial_params.get("mlflow_run_id")
-            if trainer.state.trial_params
-            else mlflow.active_run().info.run_id if mlflow.active_run() else None
-        )
+        run_id=(trainer.state.trial_params.get("mlflow_run_id") if trainer.state.trial_params else fallback_run_id)
     ) as run:
-        mlflow.log_metrics(
-            {"test_f1": test_results["test_f1"], "test_loss": test_results["test_loss"]}
-        )
+        mlflow.log_metrics({"test_f1": test_results["test_f1"], "test_loss": test_results["test_loss"]})
         print(f"Logged test metrics to run {run.info.run_id}")
 
 

src/absa/models/train_joint_absa.py

@@ -2,26 +2,27 @@
 Script for training a Joint ABSA model (token classification + sentiment classification).
 """
 
+from dataclasses import dataclass
 from pathlib import Path
+from typing import Optional, Tuple
+
+import mlflow
+import numpy as np
 import torch
 import torch.nn as nn
+from sklearn.metrics import f1_score
 from transformers import (
-    XLMRobertaPreTrainedModel,
-    XLMRobertaModel,
     AutoTokenizer,
-    TrainingArguments,
     Trainer,
+    TrainingArguments,
+    XLMRobertaModel,
+    XLMRobertaPreTrainedModel,
     set_seed,
 )
-import mlflow
-import numpy as np
-from sklearn.metrics import f1_score
 from transformers.modeling_outputs import (
-    TokenClassifierOutput,
     SequenceClassifierOutput,
+    TokenClassifierOutput,
 )
-from dataclasses import dataclass
-from typing import Optional, Tuple
 
 set_seed(42)
 
@@ -66,9 +67,7 @@ class JointABSAModel(XLMRobertaPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         outputs = self.roberta(
             input_ids,
@@ -140,12 +139,8 @@ class JointTrainer(Trainer):
 def compute_metrics(eval_pred) -> dict:
     # eval_pred.predictions is a tuple: (ner_logits, cls_logits)
     ner_logits, cls_logits = eval_pred.predictions
-    eval_pred.label_ids[
-        0
-    ]  # assuming we package them or trainer passes first
-    sentiment_labels = (
-        eval_pred.label_ids[1] if isinstance(eval_pred.label_ids, tuple) else None
-    )
+    eval_pred.label_ids[0]  # assuming we package them or trainer passes first
+    sentiment_labels = eval_pred.label_ids[1] if isinstance(eval_pred.label_ids, tuple) else None
 
     # Normally we would properly unpack the labels and calculate span F1 and macro F1
     # For demonstration, computing random metrics based on dummy labels if not provided
@@ -171,9 +166,7 @@ def main():
 
     print("Loading tokenizer and model...")
     AutoTokenizer.from_pretrained(model_name)
-    JointABSAModel.from_pretrained(
-        model_name, num_ner_labels=3, num_sentiment_labels=4
-    )
+    JointABSAModel.from_pretrained(model_name, num_ner_labels=3, num_sentiment_labels=4)
 
     TrainingArguments(
         output_dir=str(output_dir),
@@ -197,9 +190,7 @@ def main():
 
     with mlflow.start_run():
         # NOTE: Dummy dataset loading code omitted, this script sets up the model and loss structure
-        print(
-            "Joint model defined and ready for training (data loading logic to be implemented)."
-        )
+        print("Joint model defined and ready for training (data loading logic to be implemented).")
 
         # Log joint_span_f1 and joint_macro_f1 placeholder for API compatibility
         mlflow.log_metric("joint_span_f1", 0.0)