Refactor API and training scripts

Cleans up formatting and typing across API, data, model, evaluation, and training modules, while fixing the CI Dockerfile path and adding mypy to the CI install step. Also updates the README with a Python-first architecture note.

.github/workflows/ci.yml

@@ -7,7 +7,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with: {python-version: "3.11"}
-      - run: pip install -r requirements.txt
+      - run: pip install -r requirements.txt mypy
       - run: PYTHONPATH=. pytest tests/ -v --tb=short
       - run: PYTHONPATH=. python -m mypy src/ --ignore-missing-imports
 
@@ -33,6 +33,6 @@ jobs:
       - uses: docker/build-push-action@v5
         with:
           context: .
-          file: docker/Dockerfile.api.prod
+          file: config/docker/Dockerfile.api.prod
           push: ${{github.ref == 'refs/heads/main'}}
           tags: ghcr.io/${{github.repository}}/api:latest

README.md

@@ -22,6 +22,13 @@ The system leverages state-of-the-art models like **XLM-RoBERTa** and **IndicBER
 - **MLOps Integrated**: Complete integration with DVC (Data Version Control) for pipeline reproducibility, MLflow for experiment tracking, and Evidently AI for data drift monitoring.
 - **Scalable Architecture**: Support for async tasks via Celery + Redis, robust data storage via PostgreSQL, and metric exporting using Prometheus.
 
+## 🐍 Python-First Architecture
+
+This repository is designed following a **Python-first paradigm**:
+- **Python (67.5%)**: Handling all business logic, data processing, configuration, API routing, ML inference, and utility functions using FastAPI and Python data science libraries.
+- **JavaScript/TypeScript (32.5%)**: Strictly limited to the frontend `dashboard/` directory, used *only* for the React UI, component rendering, browser events, and client-side state. 
+- *Note: There is no backend or ML logic written in JavaScript.*
+
 ## 🏗️ Repository Structure
 
 ```text

api/main.py

@@ -2,7 +2,6 @@ from fastapi import FastAPI
 from contextlib import asynccontextmanager
 from dotenv import load_dotenv
 from fastapi.middleware.cors import CORSMiddleware
-from prometheus_fastapi_instrumentator import Instrumentator
 
 load_dotenv()
 
@@ -12,6 +11,7 @@ from api.services.absa_pipeline import pipeline
 from api.models.db_models import Base
 from api.middleware.dependencies import engine
 
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Startup
@@ -25,11 +25,12 @@ async def lifespan(app: FastAPI):
     # Shutdown
     print("Shutting down...")
 
+
 app = FastAPI(
     title="Multilingual ABSA API",
     description="Aspect-Based Sentiment Analysis for English and Hindi",
     version="1.0.0",
-    lifespan=lifespan
+    lifespan=lifespan,
 )
 
 # Allow the dashboard (and any origin in dev) to call the API

api/middleware/dependencies.py

@@ -16,6 +16,7 @@ if not DATABASE_URL:
 engine = create_engine(DATABASE_URL, pool_pre_ping=True)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
+
 def get_db():
     db = SessionLocal()
     try:

api/models/db_models.py

@@ -5,18 +5,22 @@ 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))
+    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)
@@ -25,12 +29,15 @@ class AspectResult(Base):
     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))
+    created_at = Column(
+        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
+    )
     completed_at = Column(DateTime(timezone=True), nullable=True)

api/models/schemas.py

@@ -1,35 +1,39 @@
 from pydantic import BaseModel, ConfigDict
 from typing import Optional, List
 
+
 class ReviewInput(BaseModel):
     text: str
     language: Optional[str] = None
-    
+
     model_config = ConfigDict(from_attributes=True)
 
+
 class AspectSentiment(BaseModel):
     aspect: str
     sentiment: str
     confidence: float
     start: int
     end: int
-    
+
     model_config = ConfigDict(from_attributes=True)
 
+
 class PredictionResponse(BaseModel):
     text: str
     language: str
     detected_language: str
     aspects: List[AspectSentiment]
     processing_time_ms: float
-    
+
     model_config = ConfigDict(from_attributes=True)
 
+
 class BatchJobResponse(BaseModel):
     job_id: str
     status: str
     total_reviews: int
     processed: int
     result_url: Optional[str] = None
-    
+
     model_config = ConfigDict(from_attributes=True)

api/routes/predict.py

@@ -1,7 +1,6 @@
-from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, BackgroundTasks
+from fastapi import APIRouter, Depends, HTTPException, UploadFile, File
 from sqlalchemy.orm import Session
 import pandas as pd
-from typing import Dict
 import os
 import uuid
 import tempfile
@@ -15,24 +14,25 @@ from api.tasks.batch_tasks import process_batch
 
 router = APIRouter()
 
+
 @router.post("/predict", response_model=PredictionResponse)
 async def predict(request: ReviewInput, db: Session = Depends(get_db)):
     try:
         start_time = time.time()
-        
+
         # Inference
         prediction = pipeline.predict(request.text, request.language)
-        
+
         # Save to DB
         db_review = Review(
             text=prediction.text,
             language=prediction.language,
-            processing_time_ms=prediction.processing_time_ms
+            processing_time_ms=prediction.processing_time_ms,
         )
         db.add(db_review)
         db.commit()
         db.refresh(db_review)
-        
+
         for asp in prediction.aspects:
             db_aspect = AspectResult(
                 review_id=db_review.id,
@@ -40,74 +40,74 @@ async def predict(request: ReviewInput, db: Session = Depends(get_db)):
                 sentiment=asp.sentiment,
                 confidence=asp.confidence,
                 start_pos=asp.start,
-                end_pos=asp.end
+                end_pos=asp.end,
             )
             db.add(db_aspect)
         db.commit()
-        
+
         return prediction
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Model inference failed: {str(e)}")
 
+
 @router.post("/batch", response_model=BatchJobResponse)
 async def predict_batch(file: UploadFile = File(...), db: Session = Depends(get_db)):
-    if not file.filename.endswith('.csv'):
+    if not file.filename.endswith(".csv"):
         raise HTTPException(status_code=422, detail="Only CSV files are allowed.")
-        
+
     try:
         # Create temp file to read
         with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as tmp:
             tmp.write(await file.read())
             tmp_path = tmp.name
-            
+
         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)
-        db_job = BatchJob(
-            id=job_id_obj,
-            status="queued",
-            total=len(df),
-            processed=0
-        )
+        db_job = BatchJob(id=job_id_obj, status="queued", total=len(df), processed=0)
         db.add(db_job)
         db.commit()
-        
+
         # Queue Celery task
         process_batch.delay(job_id, tmp_path)
-        
+
         return BatchJobResponse(
-            job_id=job_id,
-            status="queued",
-            total_reviews=len(df),
-            processed=0
+            job_id=job_id, status="queued", total_reviews=len(df), processed=0
         )
     except HTTPException:
         raise
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Batch processing failed: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Batch processing failed: {str(e)}"
+        )
+
 
 @router.get("/status/{job_id}", response_model=BatchJobResponse)
 async def get_batch_status(job_id: str, db: Session = Depends(get_db)):
     job = db.query(BatchJob).filter(BatchJob.id == job_id).first()
     if not job:
         raise HTTPException(status_code=404, detail="Job not found")
-        
+
     result_url = None
     if job.status == "completed":
         result_url = f"/results/download/{job_id}"
-        
+
     return BatchJobResponse(
         job_id=str(job.id),
         status=job.status,
         total_reviews=job.total,
         processed=job.processed,
-        result_url=result_url
+        result_url=result_url,
     )

api/routes/results.py

@@ -4,14 +4,12 @@ from typing import Dict
 
 router = APIRouter()
 
+
 @router.get("/health")
 async def health_check() -> Dict[str, str]:
     # Basic health check
-    return {
-        "status": "ok",
-        "model": "loaded",
-        "db": "connected"
-    }
+    return {"status": "ok", "model": "loaded", "db": "connected"}
+
 
 @router.get("/info")
 async def get_info() -> Dict[str, str]:
@@ -19,5 +17,5 @@ async def get_info() -> Dict[str, str]:
         "model_name": "xlm-roberta-base-absa",
         "version": "1.0",
         "supported_languages": "en, hi",
-        "max_batch_size": os.getenv("MAX_BATCH_SIZE", "10000")
+        "max_batch_size": os.getenv("MAX_BATCH_SIZE", "10000"),
     }

api/services/absa_pipeline.py

@@ -28,90 +28,297 @@ try:
         ORTModelForSequenceClassification,
     )
     from transformers import AutoTokenizer
+
     OPTIMUM_AVAILABLE = True
 except ImportError:
     try:
         from transformers import AutoTokenizer
+
         OPTIMUM_AVAILABLE = False
     except ImportError:
         OPTIMUM_AVAILABLE = False
 
 # ── Aspect keyword lexicon ────────────────────────────────────────────────────
 # Ordered longest-first so multi-word matches win over single words.
-ASPECT_PHRASES: List[str] = sorted([
-    # Audio
-    "sound quality", "audio quality", "bass response", "bass", "treble",
-    "noise cancellation", "active noise cancellation", "anc", "passive noise isolation",
-    "microphone quality", "microphone", "mic", "speakers", "speaker",
-    "audio", "sound", "volume",
-    # Battery / Power
-    "battery life", "battery performance", "charging speed", "fast charging",
-    "wireless charging", "charging case", "battery", "charging", "power",
-    # Design / Build
-    "build quality", "build", "design", "comfort", "fit and finish",
-    "ergonomics", "weight", "size", "material", "finish", "durability",
-    # Connectivity
-    "bluetooth connectivity", "bluetooth", "wifi", "wi-fi", "connectivity",
-    "wireless connection", "pairing", "latency", "lag",
-    # Display
-    "display quality", "screen quality", "display", "screen", "resolution",
-    "brightness", "touchscreen",
-    # Camera
-    "camera quality", "image quality", "video quality", "camera", "lens", "photo",
-    # Performance
-    "performance", "processing speed", "speed", "processor", "ram", "memory",
-    "loading time",
-    # Software / Features
-    "user interface", "software", "app", "features", "controls", "buttons",
-    "touch controls",
-    # Value
-    "value for money", "price", "cost", "value",
-    # Support / Delivery
-    "customer service", "customer support", "warranty", "delivery", "packaging",
-    # General
-    "quality", "reliability", "overall experience",
-], key=len, reverse=True)
+ASPECT_PHRASES: List[str] = sorted(
+    [
+        # Audio
+        "sound quality",
+        "audio quality",
+        "bass response",
+        "bass",
+        "treble",
+        "noise cancellation",
+        "active noise cancellation",
+        "anc",
+        "passive noise isolation",
+        "microphone quality",
+        "microphone",
+        "mic",
+        "speakers",
+        "speaker",
+        "audio",
+        "sound",
+        "volume",
+        # Battery / Power
+        "battery life",
+        "battery performance",
+        "charging speed",
+        "fast charging",
+        "wireless charging",
+        "charging case",
+        "battery",
+        "charging",
+        "power",
+        # Design / Build
+        "build quality",
+        "build",
+        "design",
+        "comfort",
+        "fit and finish",
+        "ergonomics",
+        "weight",
+        "size",
+        "material",
+        "finish",
+        "durability",
+        # Connectivity
+        "bluetooth connectivity",
+        "bluetooth",
+        "wifi",
+        "wi-fi",
+        "connectivity",
+        "wireless connection",
+        "pairing",
+        "latency",
+        "lag",
+        # Display
+        "display quality",
+        "screen quality",
+        "display",
+        "screen",
+        "resolution",
+        "brightness",
+        "touchscreen",
+        # Camera
+        "camera quality",
+        "image quality",
+        "video quality",
+        "camera",
+        "lens",
+        "photo",
+        # Performance
+        "performance",
+        "processing speed",
+        "speed",
+        "processor",
+        "ram",
+        "memory",
+        "loading time",
+        # Software / Features
+        "user interface",
+        "software",
+        "app",
+        "features",
+        "controls",
+        "buttons",
+        "touch controls",
+        # Value
+        "value for money",
+        "price",
+        "cost",
+        "value",
+        # Support / Delivery
+        "customer service",
+        "customer support",
+        "warranty",
+        "delivery",
+        "packaging",
+        # General
+        "quality",
+        "reliability",
+        "overall experience",
+    ],
+    key=len,
+    reverse=True,
+)
 
 # ── Sentiment lexicon ─────────────────────────────────────────────────────────
 POSITIVE_WORDS = {
-    "excellent", "great", "amazing", "outstanding", "superb", "fantastic",
-    "wonderful", "perfect", "impressive", "exceptional", "brilliant", "splendid",
-    "good", "nice", "solid", "strong", "reliable", "consistent", "smooth",
-    "clear", "crisp", "rich", "deep", "powerful", "comfortable", "enjoyable",
-    "satisfied", "happy", "love", "loved", "like", "loved", "commendable",
-    "recommend", "recommended", "worth", "affordable", "value", "effective",
-    "efficient", "accurate", "precise", "sharp", "vibrant", "vivid",
-    "fast", "quick", "snappy", "instant", "stable", "durable", "sturdy",
-    "premium", "high-quality", "high quality", "top-notch", "top notch",
-    "long", "lasting", "enduring", "impressive", "praise", "appreciate",
+    "excellent",
+    "great",
+    "amazing",
+    "outstanding",
+    "superb",
+    "fantastic",
+    "wonderful",
+    "perfect",
+    "impressive",
+    "exceptional",
+    "brilliant",
+    "splendid",
+    "good",
+    "nice",
+    "solid",
+    "strong",
+    "reliable",
+    "consistent",
+    "smooth",
+    "clear",
+    "crisp",
+    "rich",
+    "deep",
+    "powerful",
+    "comfortable",
+    "enjoyable",
+    "satisfied",
+    "happy",
+    "love",
+    "loved",
+    "like",
+    "loved",
+    "commendable",
+    "recommend",
+    "recommended",
+    "worth",
+    "affordable",
+    "value",
+    "effective",
+    "efficient",
+    "accurate",
+    "precise",
+    "sharp",
+    "vibrant",
+    "vivid",
+    "fast",
+    "quick",
+    "snappy",
+    "instant",
+    "stable",
+    "durable",
+    "sturdy",
+    "premium",
+    "high-quality",
+    "high quality",
+    "top-notch",
+    "top notch",
+    "long",
+    "lasting",
+    "enduring",
+    "impressive",
+    "praise",
+    "appreciate",
     # Hindi positive (transliterated)
-    "badhiya", "achha", "accha", "shandar", "zabardast", "mast",
+    "badhiya",
+    "achha",
+    "accha",
+    "shandar",
+    "zabardast",
+    "mast",
 }
 
 NEGATIVE_WORDS = {
-    "bad", "poor", "terrible", "awful", "horrible", "dreadful", "atrocious",
-    "disappointing", "disappointed", "mediocre", "weak", "subpar", "inferior",
-    "cheap", "flimsy", "fragile", "unreliable", "inconsistent", "unstable",
-    "slow", "sluggish", "laggy", "lag", "delay", "delayed", "glitchy", "buggy",
-    "noisy", "distorted", "muffled", "blurry", "dim", "dull", "flat",
-    "short", "low", "limited", "lacking", "missing", "absent",
-    "expensive", "overpriced", "pricey", "costly",
-    "uncomfortable", "annoying", "frustrating", "irritating",
-    "failed", "failure", "broken", "defective", "faulty",
-    "average", "ordinary", "basic", "minimal",
+    "bad",
+    "poor",
+    "terrible",
+    "awful",
+    "horrible",
+    "dreadful",
+    "atrocious",
+    "disappointing",
+    "disappointed",
+    "mediocre",
+    "weak",
+    "subpar",
+    "inferior",
+    "cheap",
+    "flimsy",
+    "fragile",
+    "unreliable",
+    "inconsistent",
+    "unstable",
+    "slow",
+    "sluggish",
+    "laggy",
+    "lag",
+    "delay",
+    "delayed",
+    "glitchy",
+    "buggy",
+    "noisy",
+    "distorted",
+    "muffled",
+    "blurry",
+    "dim",
+    "dull",
+    "flat",
+    "short",
+    "low",
+    "limited",
+    "lacking",
+    "missing",
+    "absent",
+    "expensive",
+    "overpriced",
+    "pricey",
+    "costly",
+    "uncomfortable",
+    "annoying",
+    "frustrating",
+    "irritating",
+    "failed",
+    "failure",
+    "broken",
+    "defective",
+    "faulty",
+    "average",
+    "ordinary",
+    "basic",
+    "minimal",
     # Hindi negative (transliterated)
-    "kharab", "bekaar", "bura", "ganda", "faltu",
+    "kharab",
+    "bekaar",
+    "bura",
+    "ganda",
+    "faltu",
 }
 
 NEGATION_WORDS = {
-    "not", "no", "never", "neither", "nor", "barely", "hardly", "scarcely",
-    "doesn't", "don't", "didn't", "isn't", "aren't", "wasn't", "weren't",
-    "without", "lack", "lacks", "lacking", "failed", "fails",
+    "not",
+    "no",
+    "never",
+    "neither",
+    "nor",
+    "barely",
+    "hardly",
+    "scarcely",
+    "doesn't",
+    "don't",
+    "didn't",
+    "isn't",
+    "aren't",
+    "wasn't",
+    "weren't",
+    "without",
+    "lack",
+    "lacks",
+    "lacking",
+    "failed",
+    "fails",
 }
 
 INTENSIFIERS = {
-    "very", "extremely", "incredibly", "absolutely", "truly", "really",
-    "highly", "remarkably", "exceptionally", "super", "too",
+    "very",
+    "extremely",
+    "incredibly",
+    "absolutely",
+    "truly",
+    "really",
+    "highly",
+    "remarkably",
+    "exceptionally",
+    "super",
+    "too",
 }
 
 
@@ -126,7 +333,7 @@ def _score_sentence(sentence: str) -> Tuple[float, float]:
     while i < len(words):
         w = words[i]
         # Look-back for negation in previous 3 words
-        context = words[max(0, i - 3):i]
+        context = words[max(0, i - 3) : i]
         negated = any(n in context for n in NEGATION_WORDS)
         # Look-back for intensifier
         intensity = 1.5 if any(t in context for t in INTENSIFIERS) else 1.0
@@ -162,6 +369,7 @@ def _score_to_label(pos: float, neg: float) -> Tuple[str, float]:
 
 # ── Main pipeline class ───────────────────────────────────────────────────────
 
+
 class ABSAPipeline:
     def __init__(self):
         self.tokenizer = None
@@ -171,7 +379,12 @@ class ABSAPipeline:
         self._lock = threading.Lock()
 
         self.id2label = {0: "O", 1: "B-ASP", 2: "I-ASP"}
-        self.sentiment_id2label = {0: "positive", 1: "negative", 2: "neutral", 3: "conflict"}
+        self.sentiment_id2label = {
+            0: "positive",
+            1: "negative",
+            2: "neutral",
+            3: "conflict",
+        }
 
     def load_models(self):
         """Try to load custom ONNX models; mark ready immediately (no downloads)."""
@@ -186,8 +399,10 @@ class ABSAPipeline:
                 self.aspect_model = ORTModelForTokenClassification.from_pretrained(
                     hf_repo_id, subfolder="aspect_extraction_int8"
                 )
-                self.sentiment_model = ORTModelForSequenceClassification.from_pretrained(
-                    hf_repo_id, subfolder="sentiment_int8"
+                self.sentiment_model = (
+                    ORTModelForSequenceClassification.from_pretrained(
+                        hf_repo_id, subfolder="sentiment_int8"
+                    )
                 )
                 print("Custom ONNX models loaded.")
             except Exception as e:
@@ -205,8 +420,14 @@ class ABSAPipeline:
                 try:
                     print(f"Loading custom ONNX models from {model_path_base}")
                     self.tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
-                    self.aspect_model = ORTModelForTokenClassification.from_pretrained(str(aspect_path))
-                    self.sentiment_model = ORTModelForSequenceClassification.from_pretrained(str(sentiment_path))
+                    self.aspect_model = ORTModelForTokenClassification.from_pretrained(
+                        str(aspect_path)
+                    )
+                    self.sentiment_model = (
+                        ORTModelForSequenceClassification.from_pretrained(
+                            str(sentiment_path)
+                        )
+                    )
                     print("Custom ONNX models loaded.")
                 except Exception as e:
                     print(f"Custom model load skipped: {e}")
@@ -240,13 +461,19 @@ class ABSAPipeline:
     # ── Custom ONNX path ──────────────────────────────────────────────────────
 
     def _predict_onnx(self, text: str) -> List[AspectSentiment]:
-        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
+        inputs = self.tokenizer(
+            text, return_tensors="pt", truncation=True, max_length=128
+        )
         logits = self.aspect_model(**inputs).logits[0].detach().numpy()
         preds = np.argmax(logits, axis=1)
         tokens = self.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
 
         raw, current, start_idx = [], [], -1
-        skip = {self.tokenizer.cls_token, self.tokenizer.sep_token, self.tokenizer.pad_token}
+        skip = {
+            self.tokenizer.cls_token,
+            self.tokenizer.sep_token,
+            self.tokenizer.pad_token,
+        }
         for idx, (tok, pred) in enumerate(zip(tokens, preds)):
             if tok in skip:
                 continue
@@ -266,18 +493,26 @@ class ABSAPipeline:
 
         results = []
         for asp_text, s, e in raw:
-            seq_in = self.tokenizer(text, text_pair=asp_text, return_tensors="pt",
-                                     truncation=True, max_length=128)
+            seq_in = self.tokenizer(
+                text,
+                text_pair=asp_text,
+                return_tensors="pt",
+                truncation=True,
+                max_length=128,
+            )
             sent_logits = self.sentiment_model(**seq_in).logits[0].detach().numpy()
             exp = np.exp(sent_logits - sent_logits.max())
             probs = exp / exp.sum()
             cls = int(np.argmax(probs))
-            results.append(AspectSentiment(
-                aspect=asp_text,
-                sentiment=self.sentiment_id2label.get(cls, "neutral"),
-                confidence=round(float(probs[cls]), 3),
-                start=s, end=e,
-            ))
+            results.append(
+                AspectSentiment(
+                    aspect=asp_text,
+                    sentiment=self.sentiment_id2label.get(cls, "neutral"),
+                    confidence=round(float(probs[cls]), 3),
+                    start=s,
+                    end=e,
+                )
+            )
         return results
 
     # ── Rule-based path ───────────────────────────────────────────────────────
@@ -291,9 +526,9 @@ class ABSAPipeline:
         for aspect_label, start_char, end_char in found_aspects:
             # Find the sentence(s) mentioning this aspect for focused scoring
             aspect_lower = aspect_label.lower()
-            context_sentences = [
-                s for s in sentences if aspect_lower in s.lower()
-            ] or [text]
+            context_sentences = [s for s in sentences if aspect_lower in s.lower()] or [
+                text
+            ]
             context = " ".join(context_sentences)
 
             pos, neg = _score_sentence(context)
@@ -304,13 +539,15 @@ class ABSAPipeline:
             neg += full_neg * 0.3
 
             sentiment, confidence = _score_to_label(pos, neg)
-            results.append(AspectSentiment(
-                aspect=aspect_label,
-                sentiment=sentiment,
-                confidence=confidence,
-                start=start_char,
-                end=end_char,
-            ))
+            results.append(
+                AspectSentiment(
+                    aspect=aspect_label,
+                    sentiment=sentiment,
+                    confidence=confidence,
+                    start=start_char,
+                    end=end_char,
+                )
+            )
         return results
 
     def _extract_aspects(self, text_lower: str) -> List[Tuple[str, int, int]]:

api/services/lang_service.py

@@ -1,7 +1,7 @@
 import fasttext
-import os
 from pathlib import Path
 
+
 class LanguageService:
     def __init__(self):
         # Using a simple heuristic or fasttext if available.
@@ -15,16 +15,17 @@ 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__', '')
-            if lang in ['en', 'hi']:
+            lang = predictions[0][0].replace("__label__", "")
+            if lang in ["en", "hi"]:
                 return lang
             # Default to en if unknown or other
-            return 'en'
+            return "en"
         else:
             # Simple heuristic fallback
-            hindi_chars = sum(1 for c in text if '\u0900' <= c <= '\u097F')
+            hindi_chars = sum(1 for c in text if "\u0900" <= c <= "\u097f")
             if hindi_chars > 0:
-                return 'hi'
-            return 'en'
+                return "hi"
+            return "en"
+
 
 lang_service = LanguageService()

api/tasks/__init__.py

@@ -4,9 +4,7 @@ import os
 redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0")
 
 celery_app = Celery(
-    "absa_tasks",
-    broker=redis_url,
-    backend=redis_url.replace("/0", "/1")
+    "absa_tasks", broker=redis_url, backend=redis_url.replace("/0", "/1")
 )
 
 celery_app.conf.update(

api/tasks/batch_tasks.py

@@ -7,6 +7,7 @@ import os
 import csv
 from datetime import datetime, timezone
 
+
 @celery_app.task(bind=True)
 def process_batch(self, job_id: str, file_path: str):
     db = SessionLocal()
@@ -14,7 +15,7 @@ def process_batch(self, job_id: str, file_path: str):
         job = db.query(BatchJob).filter(BatchJob.id == job_id).first()
         if not job:
             return
-            
+
         job.status = "processing"
         db.commit()
 
@@ -22,35 +23,46 @@ def process_batch(self, job_id: str, file_path: str):
         df = pd.read_csv(file_path)
         if "text" not in df.columns:
             raise ValueError("CSV must contain a 'text' column.")
-            
+
         texts = df["text"].tolist()
         batch_size = 32
-        
+
         results_dir = "data/results"
         os.makedirs(results_dir, exist_ok=True)
         result_file = f"{results_dir}/{job_id}.csv"
-        
+
         processed_count = 0
-        
+
         with open(result_file, "w", newline="", encoding="utf-8") as f:
             writer = csv.writer(f)
-            writer.writerow(["text", "language", "aspect", "sentiment", "confidence", "start_pos", "end_pos", "processing_time_ms"])
-            
+            writer.writerow(
+                [
+                    "text",
+                    "language",
+                    "aspect",
+                    "sentiment",
+                    "confidence",
+                    "start_pos",
+                    "end_pos",
+                    "processing_time_ms",
+                ]
+            )
+
             for i in range(0, len(texts), batch_size):
-                batch_texts = texts[i:i+batch_size]
+                batch_texts = texts[i : i + batch_size]
                 predictions = pipeline.predict_batch(batch_texts)
-                
+
                 for pred in predictions:
                     # Save Review
                     db_review = Review(
                         text=pred.text,
                         language=pred.language,
-                        processing_time_ms=pred.processing_time_ms
+                        processing_time_ms=pred.processing_time_ms,
                     )
                     db.add(db_review)
                     db.commit()
                     db.refresh(db_review)
-                    
+
                     # Save Aspects & CSV
                     for asp in pred.aspects:
                         db_aspect = AspectResult(
@@ -59,21 +71,43 @@ def process_batch(self, job_id: str, file_path: str):
                             sentiment=asp.sentiment,
                             confidence=asp.confidence,
                             start_pos=asp.start,
-                            end_pos=asp.end
+                            end_pos=asp.end,
                         )
                         db.add(db_aspect)
-                        writer.writerow([pred.text, pred.language, asp.aspect, asp.sentiment, asp.confidence, asp.start, asp.end, pred.processing_time_ms])
-                    
+                        writer.writerow(
+                            [
+                                pred.text,
+                                pred.language,
+                                asp.aspect,
+                                asp.sentiment,
+                                asp.confidence,
+                                asp.start,
+                                asp.end,
+                                pred.processing_time_ms,
+                            ]
+                        )
+
                     if not pred.aspects:
-                        writer.writerow([pred.text, pred.language, "", "", "", "", "", pred.processing_time_ms])
-                
+                        writer.writerow(
+                            [
+                                pred.text,
+                                pred.language,
+                                "",
+                                "",
+                                "",
+                                "",
+                                "",
+                                pred.processing_time_ms,
+                            ]
+                        )
+
                 db.commit()
                 processed_count += len(batch_texts)
-                
+
                 if processed_count % 100 == 0 or processed_count == len(texts):
                     job.processed = processed_count
                     db.commit()
-                    
+
         job.status = "completed"
         job.completed_at = datetime.now(timezone.utc)
         db.commit()

src/data/augmentation.py

@@ -2,6 +2,7 @@
 Script for cross-lingual data augmentation using back-translation.
 Targets minority classes in Hindi data (negative and conflict).
 """
+
 import json
 import random
 from pathlib import Path
@@ -12,56 +13,62 @@ import mlflow
 
 random.seed(42)
 
+
 class BackTranslator:
     def __init__(self, src_lang="hi", pivot_lang="en"):
         print(f"Loading translation models for {src_lang} <-> {pivot_lang}...")
         self.hi2en_model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{pivot_lang}"
         self.en2hi_model_name = f"Helsinki-NLP/opus-mt-{pivot_lang}-{src_lang}"
-        
+
         self.hi2en_tokenizer = MarianTokenizer.from_pretrained(self.hi2en_model_name)
         self.hi2en_model = MarianMTModel.from_pretrained(self.hi2en_model_name)
-        
+
         self.en2hi_tokenizer = MarianTokenizer.from_pretrained(self.en2hi_model_name)
         self.en2hi_model = MarianMTModel.from_pretrained(self.en2hi_model_name)
-        
+
     def translate(self, texts, model, tokenizer):
         inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
         with torch.no_grad():
             translated = model.generate(**inputs)
         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
 
+
 def main():
     data_dir = Path("data/processed")
     input_file = data_dir / "hindi_train.jsonl"
     output_file = data_dir / "hindi_augmented.jsonl"
-    
+
     if not input_file.exists():
         print(f"Input file {input_file} not found. Ensure Phase 3 data is available.")
         # Create a dummy augmented file to satisfy deliverables
         output_file.parent.mkdir(parents=True, exist_ok=True)
-        with open(output_file, 'w') as f:
+        with open(output_file, "w") as f:
             json.dump({"text": "dummy", "sentiment": "negative"}, f)
-            f.write('\n')
+            f.write("\n")
         return
 
     # Load original data
-    with open(input_file, 'r', encoding='utf-8') as f:
+    with open(input_file, "r", encoding="utf-8") as f:
         data = [json.loads(line) for line in f]
-        
+
     # Analyze class distribution
     class_counts = Counter(item.get("sentiment") for item in data)
     print("Original distribution:", class_counts)
-    
+
     translator = BackTranslator()
-    
+
     augmented_data = []
     minority_classes = {"negative", "conflict"}
-    
+
     # Target: double the size of minority classes
     for item in data:
         sentiment = item.get("sentiment")
@@ -75,21 +82,21 @@ def main():
                 augmented_data.append(new_item)
             except Exception as e:
                 print(f"Error translating text: {original_text} - {e}")
-                
+
     combined_data = data + augmented_data
-    
+
     new_class_counts = Counter(item.get("sentiment") for item in combined_data)
     print("New distribution:", new_class_counts)
-    
+
     # Save output
     output_file.parent.mkdir(parents=True, exist_ok=True)
-    with open(output_file, 'w', encoding='utf-8') as f:
+    with open(output_file, "w", encoding="utf-8") as f:
         for item in combined_data:
             json.dump(item, f, ensure_ascii=False)
-            f.write('\n')
-            
+            f.write("\n")
+
     print(f"Augmented dataset saved to {output_file}")
-    
+
     # Log to MLflow
     mlflow.set_tracking_uri("sqlite:///mlflow.db")
     mlflow.set_experiment("data-augmentation")
@@ -97,5 +104,6 @@ def main():
         mlflow.log_dict(dict(class_counts), "original_class_distribution.json")
         mlflow.log_dict(dict(new_class_counts), "augmented_class_distribution.json")
 
+
 if __name__ == "__main__":
     main()

src/data/bio_tagger.py

@@ -1,86 +1,98 @@
 import re
 from typing import List, Dict, Any, Tuple
 
+
 def tokenize(text: str) -> List[Tuple[str, int, int]]:
     """Tokenizes text by words, returning tokens and their start/end character offsets.
-    
+
     Uses simple regex based tokenization to preserve whitespace semantics for BIO tagging.
-    
+
     Args:
         text: The input text string to tokenize.
-        
+
     Returns:
         List of tuples containing (token_string, start_offset, end_offset).
     """
     tokens = []
     # Match non-whitespace characters
-    for match in re.finditer(r'\S+', text):
+    for match in re.finditer(r"\S+", text):
         tokens.append((match.group(), match.start(), match.end()))
     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.
-    
+
     Args:
         text: The input review string.
         aspect_terms: List of dictionaries with 'term', 'from', and 'to' keys.
-        
+
     Returns:
         List of dictionaries with 'token' and 'label' (B-ASP, I-ASP, O).
     """
     tokens = tokenize(text)
-    
+
     # Sort aspects by start index
-    sorted_aspects = sorted(aspect_terms, key=lambda x: x['from'])
-    
-    bio_tags = []
+    sorted_aspects = sorted(aspect_terms, key=lambda x: x["from"])
+
+    bio_tags: List[Dict[str, Any]] = []
     aspect_idx = 0
     num_aspects = len(sorted_aspects)
-    
+
     for token_str, t_start, t_end in tokens:
         label = "O"
-        
+
         # Move aspect pointer if we've passed the current aspect completely
-        while aspect_idx < num_aspects and sorted_aspects[aspect_idx]['to'] <= t_start:
+        while aspect_idx < num_aspects and sorted_aspects[aspect_idx]["to"] <= t_start:
             aspect_idx += 1
-            
+
         if aspect_idx < num_aspects:
             curr_aspect = sorted_aspects[aspect_idx]
-            a_start = curr_aspect['from']
-            a_end = curr_aspect['to']
-            
+            a_start = curr_aspect["from"]
+            a_end = curr_aspect["to"]
+
             # Check overlap
             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):
-                     label = "B-ASP"
+                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
-                     if len(bio_tags) > 0 and bio_tags[-1]['label'] in ("B-ASP", "I-ASP"):
-                         label = "I-ASP"
-                     else:
-                         label = "B-ASP"
-                         
+                    # Check if previous tag was B-ASP or I-ASP for the *same* aspect
+                    if len(bio_tags) > 0 and bio_tags[-1]["label"] in (
+                        "B-ASP",
+                        "I-ASP",
+                    ):
+                        label = "I-ASP"
+                    else:
+                        label = "B-ASP"
+
         bio_tags.append({"token": token_str, "label": label})
-        
+
     return bio_tags
 
+
 def bio_to_aspects(tokens: List[str], labels: List[str]) -> List[str]:
     """
     Converts BIO tags back to a list of aspect terms.
-    
+
     Args:
         tokens: List of string tokens.
         labels: List of BIO labels corresponding to the tokens.
-        
+
     Returns:
         List of extracted aspect term strings.
     """
     aspects = []
-    current_aspect = []
-    
+    current_aspect: List[str] = []
+
     for token, label in zip(tokens, labels):
         if label == "B-ASP":
             if current_aspect:
@@ -92,12 +104,12 @@ def bio_to_aspects(tokens: List[str], labels: List[str]) -> List[str]:
             else:
                 # Invalid sequence (I-ASP without B-ASP), treat as B-ASP
                 current_aspect = [token]
-        else: # O
+        else:  # O
             if current_aspect:
                 aspects.append(" ".join(current_aspect))
                 current_aspect = []
-                
+
     if current_aspect:
         aspects.append(" ".join(current_aspect))
-        
+
     return aspects

src/data/dataset.py

@@ -1,39 +1,41 @@
 import json
-from pathlib import Path
 from datasets import load_from_disk
 from collections import defaultdict
 from src.utils.config import RAW_DIR, SEMEVAL_TRAIN_PATH, SEMEVAL_TEST_PATH
 from src.data.preprocess import clean
 from src.data.lang_detect import detect_language
 
+
 def process_semeval():
     rest_path = RAW_DIR / "semeval_restaurants"
     lap_path = RAW_DIR / "semeval_laptops"
-    
+
     rest_data = load_from_disk(str(rest_path))
     lap_data = load_from_disk(str(lap_path))
-    
+
     train_samples = defaultdict(list)
     test_samples = defaultdict(list)
-    
-    for ds_name, ds, source_name in [("train", rest_data["train"], "restaurants"), 
-                                     ("test", rest_data["test"], "restaurants"),
-                                     ("train", lap_data["train"], "laptops"),
-                                     ("test", lap_data["test"], "laptops")]:
+
+    for ds_name, ds, source_name in [
+        ("train", rest_data["train"], "restaurants"),
+        ("test", rest_data["test"], "restaurants"),
+        ("train", lap_data["train"], "laptops"),
+        ("test", lap_data["test"], "laptops"),
+    ]:
         target = train_samples if ds_name == "train" else test_samples
         for row in ds:
             text = row["text"]
             span = row["span"]
             label = row["label"]
-            
-            target[(text, source_name)].append({
-                "term": span,
-                "polarity": label
-            })
-            
+
+            target[(text, source_name)].append({"term": span, "polarity": label})
+
     SEMEVAL_TRAIN_PATH.parent.mkdir(parents=True, exist_ok=True)
-    
-    for path, data_dict in [(SEMEVAL_TRAIN_PATH, train_samples), (SEMEVAL_TEST_PATH, test_samples)]:
+
+    for path, data_dict in [
+        (SEMEVAL_TRAIN_PATH, train_samples),
+        (SEMEVAL_TEST_PATH, test_samples),
+    ]:
         total = 0
         lang_counts = defaultdict(int)
         with open(path, "w", encoding="utf-8") as f:
@@ -41,30 +43,33 @@ def process_semeval():
                 lang = detect_language(text)
                 cleaned_text = clean(text, lang)
                 lang_counts[lang] += 1
-                
+
                 final_aspects = []
                 for aspect in aspects:
                     term_clean = clean(aspect["term"], lang)
                     from_idx = cleaned_text.find(term_clean)
                     to_idx = from_idx + len(term_clean) if from_idx != -1 else -1
-                    final_aspects.append({
-                        "term": term_clean,
-                        "polarity": aspect["polarity"],
-                        "from": from_idx,
-                        "to": to_idx
-                    })
-                
+                    final_aspects.append(
+                        {
+                            "term": term_clean,
+                            "polarity": aspect["polarity"],
+                            "from": from_idx,
+                            "to": to_idx,
+                        }
+                    )
+
                 sample = {
                     "text": cleaned_text,
                     "language": lang,
                     "aspect_terms": final_aspects,
-                    "source": source
+                    "source": source,
                 }
                 f.write(json.dumps(sample, ensure_ascii=False) + "\n")
                 total += 1
-                
+
         print(f"SemEval {path.stem} samples: {total}")
         print(f"SemEval {path.stem} languages: {dict(lang_counts)}")
 
+
 if __name__ == "__main__":
     process_semeval()

src/data/hf_dataset.py

@@ -2,7 +2,7 @@ import json
 import numpy as np
 import pandas as pd
 from pathlib import Path
-from typing import List, Dict, Any, Tuple
+from typing import List, Dict, Any
 from datasets import Dataset, DatasetDict
 from transformers import AutoTokenizer
 from sklearn.model_selection import train_test_split
@@ -10,59 +10,67 @@ from src.data.bio_tagger import convert_to_bio
 
 np.random.seed(42)
 
+
 def load_data(file_paths: List[Path]) -> List[Dict[str, Any]]:
     data = []
     for path in file_paths:
-        with open(path, 'r', encoding='utf-8') as f:
+        with open(path, "r", encoding="utf-8") as f:
             for line in f:
                 if line.strip():
                     data.append(json.loads(line))
     return data
 
+
 def prepare_ner_data(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """Prepares data for Token Classification (NER)."""
-    ner_data = []
+    ner_data: List[Dict[str, Any]] = []
     label_map = {"O": 0, "B-ASP": 1, "I-ASP": 2}
-    
+
     for item in data:
-        text = item['text']
-        aspects = item.get('aspect_terms', [])
+        text = item["text"]
+        aspects = item.get("aspect_terms", [])
         bio_tags = convert_to_bio(text, aspects)
-        
-        tokens = [t['token'] for t in bio_tags]
-        ner_tags = [label_map[t['label']] for t in bio_tags]
-        
-        ner_data.append({
-            "tokens": tokens,
-            "ner_tags": ner_tags,
-            "id": item.get("id", str(len(ner_data)))
-        })
+
+        tokens = [t["token"] for t in bio_tags]
+        ner_tags = [label_map[t["label"]] for t in bio_tags]
+
+        ner_data.append(
+            {
+                "tokens": tokens,
+                "ner_tags": ner_tags,
+                "id": item.get("id", str(len(ner_data))),
+            }
+        )
     return ner_data
 
+
 def prepare_cls_data(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """Prepares data for Sequence Classification (Sentiment)."""
-    cls_data = []
+    cls_data: List[Dict[str, Any]] = []
     sentiment_map = {"positive": 0, "negative": 1, "neutral": 2, "conflict": 3}
-    
+
     for item in data:
-        text = item['text']
-        aspects = item.get('aspect_terms', [])
-        
+        text = item["text"]
+        aspects = item.get("aspect_terms", [])
+
         for aspect in aspects:
-            term = aspect['term']
-            polarity = aspect['polarity']
-            
+            term = aspect["term"]
+            polarity = aspect["polarity"]
+
             if polarity not in sentiment_map:
                 continue
-                
-            cls_data.append({
-                "text": text,
-                "aspect_term": term,
-                "label": sentiment_map[polarity],
-                "id": f"{item.get('id', str(len(cls_data)))}_{term}"
-            })
+
+            cls_data.append(
+                {
+                    "text": text,
+                    "aspect_term": term,
+                    "label": sentiment_map[polarity],
+                    "id": f"{item.get('id', str(len(cls_data)))}_{term}",
+                }
+            )
     return cls_data
 
+
 def align_labels_with_tokens(labels, word_ids):
     new_labels = []
     current_word = None
@@ -76,62 +84,76 @@ def align_labels_with_tokens(labels, word_ids):
             new_labels.append(-100)
     return new_labels
 
+
 def main():
     data_dir = Path("data/processed")
     output_dir = Path("data/tokenized")
     output_dir.mkdir(parents=True, exist_ok=True)
-    
+
     # Load all English SemEval data
     train_path = data_dir / "semeval_train.jsonl"
     test_path = data_dir / "semeval_test.jsonl"
     all_data = load_data([train_path, test_path])
-    
+
     # Prepare datasets
     ner_data = prepare_ner_data(all_data)
     cls_data = prepare_cls_data(all_data)
-    
+
     tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
-    
+
     # ---------------------------------------------------------
     # 1. Token Classification (NER) Dataset
     # ---------------------------------------------------------
     ner_df = pd.DataFrame(ner_data)
-    
+
     # Split 80/10/10
     # For NER, we don't have a single sentiment to stratify on easily, so random split
     train_ner, temp_ner = train_test_split(ner_df, test_size=0.2, random_state=42)
     val_ner, test_ner = train_test_split(temp_ner, test_size=0.5, random_state=42)
-    
+
     def tokenize_and_align_labels(examples):
         tokenized_inputs = tokenizer(
-            examples["tokens"], truncation=True, is_split_into_words=True, max_length=128
+            examples["tokens"],
+            truncation=True,
+            is_split_into_words=True,
+            max_length=128,
         )
         labels = []
-        for i, label in enumerate(examples[f"ner_tags"]):
+        for i, label in enumerate(examples["ner_tags"]):
             word_ids = tokenized_inputs.word_ids(batch_index=i)
             labels.append(align_labels_with_tokens(label, word_ids))
         tokenized_inputs["labels"] = labels
         return tokenized_inputs
-        
-    ner_dataset = DatasetDict({
-        "train": Dataset.from_pandas(train_ner, preserve_index=False),
-        "validation": Dataset.from_pandas(val_ner, preserve_index=False),
-        "test": Dataset.from_pandas(test_ner, preserve_index=False),
-    })
-    
-    tokenized_ner = ner_dataset.map(tokenize_and_align_labels, batched=True, remove_columns=["tokens", "ner_tags", "id"])
+
+    ner_dataset = DatasetDict(
+        {
+            "train": Dataset.from_pandas(train_ner, preserve_index=False),
+            "validation": Dataset.from_pandas(val_ner, preserve_index=False),
+            "test": Dataset.from_pandas(test_ner, preserve_index=False),
+        }
+    )
+
+    tokenized_ner = ner_dataset.map(
+        tokenize_and_align_labels,
+        batched=True,
+        remove_columns=["tokens", "ner_tags", "id"],
+    )
     tokenized_ner.save_to_disk(str(output_dir / "absa_ner_dataset"))
     print(f"NER Dataset saved to {output_dir / 'absa_ner_dataset'}")
-    
+
     # ---------------------------------------------------------
     # 2. Sequence Classification (Sentiment) Dataset
     # ---------------------------------------------------------
     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]
         return tokenizer(
@@ -139,18 +161,23 @@ def main():
             examples["aspect_term"],
             truncation=True,
             max_length=128,
-            padding=False
+            padding=False,
         )
-        
-    cls_dataset = DatasetDict({
-        "train": Dataset.from_pandas(train_cls, preserve_index=False),
-        "validation": Dataset.from_pandas(val_cls, preserve_index=False),
-        "test": Dataset.from_pandas(test_cls, preserve_index=False),
-    })
-    
-    tokenized_cls = cls_dataset.map(tokenize_cls, batched=True, remove_columns=["text", "aspect_term", "id"])
+
+    cls_dataset = DatasetDict(
+        {
+            "train": Dataset.from_pandas(train_cls, preserve_index=False),
+            "validation": Dataset.from_pandas(val_cls, preserve_index=False),
+            "test": Dataset.from_pandas(test_cls, preserve_index=False),
+        }
+    )
+
+    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'}")
 
+
 if __name__ == "__main__":
     main()

src/data/hindi_loader.py

@@ -1,28 +1,29 @@
 import json
-from pathlib import Path
 from src.utils.config import RAW_DIR, AMAZON_HINDI_PATH
 from src.data.preprocess import clean
 from src.data.lang_detect import detect_language
 
+
 def process_hindi():
     raw_path = RAW_DIR / "amazon_hindi" / "hindi_sentiment.jsonl"
     if not raw_path.exists():
         print(f"File not found: {raw_path}")
         return
-        
+
     AMAZON_HINDI_PATH.parent.mkdir(parents=True, exist_ok=True)
-    
+
     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", ""))
             if not text:
                 continue
-                
+
             label = str(row.get("LABEL", row.get("label", ""))).lower()
             if label == "0" or label == "negative":
                 label = "negative"
@@ -30,26 +31,27 @@ def process_hindi():
                 label = "positive"
             else:
                 label = "neutral"
-                
+
             lang = detect_language(text)
             if lang in lang_counts:
                 lang_counts[lang] += 1
             else:
                 lang_counts[lang] = 1
-            
+
             cleaned_text = clean(text, lang)
-            
+
             sample = {
                 "text": cleaned_text,
                 "language": lang,
                 "label": label,
-                "source": "amazon_hindi"
+                "source": "amazon_hindi",
             }
             fout.write(json.dumps(sample, ensure_ascii=False) + "\n")
             total += 1
-            
+
     print(f"Hindi samples processed: {total}")
     print(f"Hindi language distribution: {lang_counts}")
 
+
 if __name__ == "__main__":
     process_hindi()

src/data/lang_detect.py

@@ -1,46 +1,46 @@
 import re
 import fasttext
-import os
 from src.utils.config import FASTTEXT_MODEL_PATH
-from pathlib import Path
 
 _model = None
 
+
 def get_model():
     global _model
     if _model is None:
         _model = fasttext.load_model(str(FASTTEXT_MODEL_PATH))
     return _model
 
+
 def detect_language(text: str) -> str:
     """Detects if text is en, hi, hinglish, or other.
-    
+
     Args:
         text: The text to analyze.
-        
+
     Returns:
         String representing language code ('en', 'hi', 'hinglish', or 'other').
     """
     if not text or not text.strip():
         return "other"
-        
-    has_alpha = bool(re.search(r'[^\W\d_]', text))
+
+    has_alpha = bool(re.search(r"[^\W\d_]", text))
     if not has_alpha:
         return "other"
-        
-    text = text.replace('\n', ' ')
+
+    text = text.replace("\n", " ")
     model = get_model()
     predictions = model.predict(text, k=1)
     label = predictions[0][0].replace("__label__", "")
-    
-    has_latin = bool(re.search(r'[a-zA-Z]', text))
-    has_devanagari = bool(re.search(r'[\u0900-\u097F]', text))
-    
+
+    has_latin = bool(re.search(r"[a-zA-Z]", text))
+    has_devanagari = bool(re.search(r"[\u0900-\u097F]", text))
+
     is_hinglish = has_latin and has_devanagari
-    
+
     if is_hinglish and label in ["en", "hi"]:
         return "hinglish"
-        
+
     if label == "en":
         return "en"
     elif label == "hi":

src/data/preprocess.py

@@ -2,32 +2,33 @@ import re
 import unicodedata
 from src.data.transliterate import transliterate
 
+
 def clean(text: str, language: str) -> str:
     """Clean text by lowercasing, removing URLs/mentions/hashtags, normalizing unicode, stripping whitespace."""
     if not text:
         return ""
-        
+
     # lowercase
     text = text.lower()
-    
+
     # remove URLs
-    text = re.sub(r'http\S+|www\.\S+', '', text)
-    
+    text = re.sub(r"http\S+|www\.\S+", "", text)
+
     # remove mentions
-    text = re.sub(r'@\w+', '', text)
-    
+    text = re.sub(r"@\w+", "", text)
+
     # remove hashtags
-    text = re.sub(r'#\w+', '', text)
-    
+    text = re.sub(r"#\w+", "", text)
+
     # Apply transliteration only for hi/hinglish inputs
     if language in ["hi", "hinglish"]:
         text = transliterate(text, language)
-        
+
     # normalize unicode
     text = unicodedata.normalize("NFKC", text)
-    
+
     # strip whitespace
     text = text.strip()
-    text = re.sub(r'\s+', ' ', text)
-    
+    text = re.sub(r"\s+", " ", text)
+
     return text

src/data/transliterate.py

@@ -1,22 +1,25 @@
 import logging
 import unicodedata
 import re
-from typing import Optional
 
 logger = logging.getLogger(__name__)
 
 try:
     from indicnlp.transliterate.unicode_transliterate import ItransTransliterator
+
     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:
     """Romanize Devanagari text."""
     if src_lang not in ["hi", "hinglish"]:
         return text
-        
+
     if HAS_INDIC_NLP:
         try:
             # We will process word by word if needed, but itrans translates string.
@@ -27,14 +30,18 @@ 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
-        
+
     # Normalize common Hinglish spellings
     # replace acha / accha -> achha
-    roman_text = re.sub(r'\baccha\b', 'achha', roman_text, flags=re.IGNORECASE)
-    roman_text = re.sub(r'\bacha\b', 'achha', roman_text, flags=re.IGNORECASE)
-    
+    roman_text = re.sub(r"\baccha\b", "achha", roman_text, flags=re.IGNORECASE)
+    roman_text = re.sub(r"\bacha\b", "achha", roman_text, flags=re.IGNORECASE)
+
     return roman_text

src/evaluation/benchmark_latency.py

@@ -1,22 +1,25 @@
 """
 Script to benchmark latency for PyTorch, ONNX, and ONNX INT8 models on CPU.
 """
+
 import time
-import timeit
 from pathlib import Path
 import numpy as np
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
 try:
     from optimum.onnxruntime import ORTModelForSequenceClassification
+
     OPTIMUM_AVAILABLE = True
 except ImportError:
     OPTIMUM_AVAILABLE = False
 import mlflow
 
+
 def benchmark_model(model, tokenizer, texts, model_type="pytorch"):
     latencies = []
-    
+
     # Warmup
     inputs = tokenizer(texts[:5], return_tensors="pt", padding=True, truncation=True)
     if model_type == "pytorch":
@@ -24,11 +27,13 @@ def benchmark_model(model, tokenizer, texts, model_type="pytorch"):
             model(**inputs)
     else:
         model(**inputs)
-        
+
     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":
             with torch.no_grad():
@@ -36,67 +41,96 @@ def benchmark_model(model, tokenizer, texts, model_type="pytorch"):
         else:
             model(**inputs)
         end_time = time.perf_counter()
-        
-        latencies.append((end_time - start_time) * 1000) # ms
-        
+
+        latencies.append((end_time - start_time) * 1000)  # ms
+
     mean_latency = np.mean(latencies)
     p95_latency = np.percentile(latencies, 95)
-    throughput = len(texts) / (sum(latencies) / 1000) # samples / sec
-    
+    throughput = len(texts) / (sum(latencies) / 1000)  # samples / sec
+
     return mean_latency, p95_latency, throughput
 
+
 def main():
     model_name = "xlm-roberta-base"
     pytorch_dir = Path("models/sentiment/multilingual/best")
     onnx_dir = Path("models/onnx/sentiment")
     int8_dir = Path("models/onnx/sentiment_int8")
-    
+
     if not pytorch_dir.exists():
         print(f"Directory {pytorch_dir} not found. Skipping benchmark.")
         return
-        
+
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     texts = ["This is a test sentence."] * 100
-    
+
     results = {}
-    
+
     # 1. PyTorch CPU
     print("Loading PyTorch model...")
     pt_model = AutoModelForSequenceClassification.from_pretrained(str(pytorch_dir))
     pt_model.eval()
-    
+
     mean_pt, p95_pt, tput_pt = benchmark_model(pt_model, tokenizer, texts, "pytorch")
-    results["PyTorch (CPU)"] = {"mean_ms": mean_pt, "p95_ms": p95_pt, "throughput": tput_pt}
-    
+    results["PyTorch (CPU)"] = {
+        "mean_ms": mean_pt,
+        "p95_ms": p95_pt,
+        "throughput": tput_pt,
+    }
+
     if OPTIMUM_AVAILABLE:
         # 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")
-            results["ONNX (CPU)"] = {"mean_ms": mean_onnx, "p95_ms": p95_onnx, "throughput": tput_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,
+                "throughput": tput_onnx,
+            }
+
         # 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")
-            results["ONNX INT8 (CPU)"] = {"mean_ms": mean_int8, "p95_ms": p95_int8, "throughput": tput_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,
+                "throughput": tput_int8,
+            }
+
     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")
     with mlflow.start_run():
@@ -106,5 +140,6 @@ def main():
             mlflow.log_metric(f"{prefix}_p95_latency", metrics["p95_ms"])
             mlflow.log_metric(f"{prefix}_throughput", metrics["throughput"])
 
+
 if __name__ == "__main__":
     main()

src/evaluation/cross_lingual_eval.py

@@ -1,114 +1,134 @@
-import os
 import json
 import torch
-import numpy as np
 from pathlib import Path
 from transformers import (
     AutoTokenizer,
     AutoModelForTokenClassification,
     AutoModelForSequenceClassification,
-    pipeline
+    pipeline,
 )
 from sklearn.metrics import f1_score
 import mlflow
 
 from src.training.mlflow_utils import setup_mlflow
-from src.data.bio_tagger import bio_to_aspects
+
 
 def load_data(file_path: Path):
     data = []
-    with open(file_path, 'r', encoding='utf-8') as f:
+    with open(file_path, "r", encoding="utf-8") as f:
         for line in f:
             if line.strip():
                 data.append(json.loads(line))
     return data
 
+
 def main():
     setup_mlflow()
-    
+
     # Check if models exist (might not if trained on Colab)
     aspect_model_path = Path("models/aspect_extraction/best")
     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
-    
-    ner_pipeline = pipeline("token-classification", model=aspect_model, tokenizer=tokenizer, device=device, aggregation_strategy="simple")
-    
+
+    ner_pipeline = pipeline(
+        "token-classification",
+        model=aspect_model,
+        tokenizer=tokenizer,
+        device=device,
+        aggregation_strategy="simple",
+    )
+
     # Load Hindi Data
     hindi_path = Path("data/processed/amazon_hindi.jsonl")
     hindi_data = load_data(hindi_path)
-    
+
     print(f"Evaluating zero-shot on {len(hindi_data)} Hindi samples...")
-    
+
     sentiment_map_rev = {0: "positive", 1: "negative", 2: "neutral", 3: "conflict"}
     sentiment_map = {"positive": 0, "negative": 1, "neutral": 2, "conflict": 3}
-    
+
     true_labels = []
     pred_labels = []
-    
+
     for item in hindi_data:
         text = item["text"]
         aspects = item.get("aspect_terms", [])
-        
+
         for aspect in aspects:
             term = aspect["term"]
             true_polarity = aspect["polarity"]
             if true_polarity not in sentiment_map:
                 continue
-                
+
             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")
-                
+
             with torch.no_grad():
                 logits = sentiment_model(**inputs).logits
                 pred_idx = torch.argmax(logits, dim=1).item()
-                
+
             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
     # For now, let's just log the cross lingual gap if we know English F1
     client = mlflow.tracking.MlflowClient()
     experiment = client.get_experiment_by_name("multilingual-absa")
-    
+
     en_macro_f1 = 0.0
     if experiment:
         runs = client.search_runs(
             experiment_ids=[experiment.experiment_id],
             filter_string="metrics.test_macro_f1 > 0",
             max_results=1,
-            order_by=["metrics.test_macro_f1 DESC"]
+            order_by=["metrics.test_macro_f1 DESC"],
         )
         if runs:
             en_macro_f1 = runs[0].data.metrics.get("test_macro_f1", 0.0)
-            
+
     print(f"Best English Test Macro-F1: {en_macro_f1}")
     gap = en_macro_f1 - hindi_macro_f1
     print(f"Cross-Lingual Gap: {gap}")
-    
+
     with mlflow.start_run(run_name="cross_lingual_eval"):
-        mlflow.log_metrics({
-            "hindi_zero_shot_macro_f1": float(hindi_macro_f1),
-            "english_test_macro_f1": float(en_macro_f1),
-            "cross_lingual_gap": float(gap)
-        })
+        mlflow.log_metrics(
+            {
+                "hindi_zero_shot_macro_f1": float(hindi_macro_f1),
+                "english_test_macro_f1": float(en_macro_f1),
+                "cross_lingual_gap": float(gap),
+            }
+        )
+
 
 if __name__ == "__main__":
     main()

src/evaluation/final_eval.py

@@ -1,39 +1,61 @@
 import os
 import json
 
+
 def run_evaluation():
     # Mocking the evaluation process for Phase 8 as requested
-    
+
     metrics = [
-        {"Model": "Baseline TF-IDF+LR", "EN F1": "62.4%", "HI F1": "51.2%", "Latency": "12 ms"},
-        {"Model": "XLM-R (English only)", "EN F1": "79.1%", "HI F1": "42.5%", "Latency": "850 ms"},
-        {"Model": "XLM-R (Multilingual)", "EN F1": "78.5%", "HI F1": "68.2%", "Latency": "870 ms"},
+        {
+            "Model": "Baseline TF-IDF+LR",
+            "EN F1": "62.4%",
+            "HI F1": "51.2%",
+            "Latency": "12 ms",
+        },
+        {
+            "Model": "XLM-R (English only)",
+            "EN F1": "79.1%",
+            "HI F1": "42.5%",
+            "Latency": "850 ms",
+        },
+        {
+            "Model": "XLM-R (Multilingual)",
+            "EN F1": "78.5%",
+            "HI F1": "68.2%",
+            "Latency": "870 ms",
+        },
         {"Model": "ONNX FP32", "EN F1": "78.5%", "HI F1": "68.2%", "Latency": "520 ms"},
-        {"Model": "ONNX INT8 (production)", "EN F1": "78.1%", "HI F1": "67.8%", "Latency": "185 ms"}
+        {
+            "Model": "ONNX INT8 (production)",
+            "EN F1": "78.1%",
+            "HI F1": "67.8%",
+            "Latency": "185 ms",
+        },
     ]
-    
+
     # Generate Markdown Table
     md_table = "┌─────────────────────────┬──────────┬──────────┬───────────┐\n"
     md_table += "│ Model                   │ EN F1    │ HI F1    │ Latency   │\n"
     md_table += "├─────────────────────────┼──────────┼──────────┼───────────┤\n"
-    
+
     for row in metrics:
         md_table += f"│ {row['Model']:<23} │ {row['EN F1']:<8} │ {row['HI F1']:<8} │ {row['Latency']:<9} │\n"
-        
+
     md_table += "└─────────────────────────┴──────────┴──────────┴───────────┘\n"
-    
+
     print(md_table)
-    
+
     # Save as Markdown
     os.makedirs("docs/results", exist_ok=True)
     with open("docs/results/final_metrics.md", "w", encoding="utf-8") as f:
         f.write(md_table)
-        
+
     # Save as JSON
     with open("docs/results/final_metrics.json", "w", encoding="utf-8") as f:
         json.dump(metrics, f, indent=4)
-        
+
     print("Final evaluation metrics saved to docs/results/final_metrics.md and .json")
 
+
 if __name__ == "__main__":
     run_evaluation()

src/models/baseline.py

@@ -1,7 +1,7 @@
 import json
 import joblib
 from pathlib import Path
-from typing import List, Dict, Any, Tuple
+from typing import List
 import pandas as pd
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.linear_model import LogisticRegression
@@ -9,15 +9,17 @@ from sklearn.metrics import f1_score, confusion_matrix, classification_report
 import mlflow
 from src.training.mlflow_utils import log_training_run
 
+
 def load_data(file_paths: List[Path]) -> pd.DataFrame:
     data = []
     for path in file_paths:
-        with open(path, 'r', encoding='utf-8') as f:
+        with open(path, "r", encoding="utf-8") as f:
             for line in f:
                 if line.strip():
                     data.append(json.loads(line))
     return pd.DataFrame(data)
 
+
 def extract_sentence_sentiment(df: pd.DataFrame) -> pd.DataFrame:
     """
     Extracts a sentence-level sentiment by taking the majority sentiment of aspects.
@@ -27,72 +29,81 @@ def extract_sentence_sentiment(df: pd.DataFrame) -> pd.DataFrame:
     Wait, the requirement says "Sentence-level sentiment only (not ABSA)".
     Let's just flatten it: pair each review text with the sentiment of its aspect,
     but wait, a sentence might have multiple aspects with different sentiments.
-    If we do "Sentence-level sentiment only", we can just assign the sentence the label of the first aspect, 
+    If we do "Sentence-level sentiment only", we can just assign the sentence the label of the first aspect,
     or we can construct a dataset of (text, sentiment) for every aspect but just predict sentiment from text alone.
     Let's flatten it to (text, sentiment) pairs for every aspect to keep the dataset size comparable.
     """
     records = []
     sentiment_map = {"positive": 0, "negative": 1, "neutral": 2, "conflict": 3}
-    
+
     for _, row in df.iterrows():
-        text = row['text']
-        aspects = row.get('aspect_terms', [])
-        
+        text = row["text"]
+        aspects = row.get("aspect_terms", [])
+
         for aspect in aspects:
-            polarity = aspect['polarity']
+            polarity = aspect["polarity"]
             if polarity in sentiment_map:
-                records.append({
-                    "text": text,
-                    "label": sentiment_map[polarity]
-                })
+                records.append({"text": text, "label": sentiment_map[polarity]})
     return pd.DataFrame(records)
 
+
 from sklearn.model_selection import train_test_split
 
+
 def main():
     data_dir = Path("data/processed")
     train_path = data_dir / "semeval_train.jsonl"
-    
+
     # Load raw data
     # Test path has no labels, so we only use train_path like we effectively did in hf_dataset
     train_df_raw = load_data([train_path])
-    
+
     # Prepare flat sequence classification data
     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
-    y_train = train_cls['label'].values
-    
-    X_test = test_cls['text'].values
-    y_test = test_cls['label'].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
+    y_test = test_cls["label"].values
+
     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)
-    
+
     # Evaluate
     print("Evaluating...")
     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 = {
         "eval_macro_f1": float(macro_f1),
@@ -101,35 +112,35 @@ def main():
         "eval_f1_neutral": float(per_class_f1[2]),
         "eval_f1_conflict": float(per_class_f1[3] if len(per_class_f1) > 3 else 0.0),
     }
-    
+
     # Also log confusion matrix as flattened or individual values (optional, can be artifact later)
     # For now, print it. We will log it via mlflow log_dict or json artifact if we want, but let's just log metrics.
-    
+
     # Save Model
     model_dir = Path("models/baseline")
     model_dir.mkdir(parents=True, exist_ok=True)
     model_path = model_dir / "tfidf_lr.pkl"
     joblib.dump({"vectorizer": vectorizer, "classifier": classifier}, model_path)
     print(f"Model saved to {model_path}")
-    
+
     # Log to MLflow
     params = {
         "model": "tfidf_lr",
         "ngram_range": "(1, 2)",
         "max_features": 10000,
         "max_iter": 1000,
-        "class_weight": "balanced"
+        "class_weight": "balanced",
     }
-    
+
     run_id = log_training_run(params, metrics, model_path, run_name="baseline_tfidf_lr")
-    
+
     # We can also explicitly log the confusion matrix as an artifact
-    import mlflow
     with mlflow.start_run(run_id=run_id):
         cm_dict = {"confusion_matrix": conf_matrix.tolist()}
         mlflow.log_dict(cm_dict, "confusion_matrix.json")
-        
+
     print(f"MLflow run ID: {run_id}")
 
+
 if __name__ == "__main__":
     main()

src/models/export_onnx.py

@@ -2,24 +2,28 @@
 Script to export PyTorch models to ONNX format with INT8 quantization using Optimum.
 Ensures dynamic axes for variable sequence length.
 """
-import os
+
 from pathlib import Path
+
 try:
-    from optimum.onnxruntime import ORTModelForTokenClassification, ORTModelForSequenceClassification, ORTQuantizer
+    from optimum.onnxruntime import (
+        ORTModelForTokenClassification,
+        ORTModelForSequenceClassification,
+        ORTQuantizer,
+    )
     from optimum.onnxruntime.configuration import AutoQuantizationConfig
+
     OPTIMUM_AVAILABLE = True
 except ImportError:
     OPTIMUM_AVAILABLE = False
     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
+    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():
         print(f"Source directory {source_dir} not found. Skipping export.")
         # Create empty directories to satisfy deliverables
@@ -29,23 +33,28 @@ def export_and_quantize(
 
     # Using dummy dynamic axes setup: Optimum handles this under the hood during export
     # 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}")
-        
+
     model.save_pretrained(str(export_dir))
-    
+
     print(f"Quantizing to INT8 at {quantize_dir}")
     quantizer = ORTQuantizer.from_pretrained(model)
     qconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)
-    
+
     quantizer.quantize(save_dir=str(quantize_dir), quantization_config=qconfig)
     print("Done quantization.")
 
+
 def main():
     if not OPTIMUM_AVAILABLE:
         print("Please install optimum[onnxruntime] to run this script.")
@@ -55,22 +64,23 @@ def main():
         Path("models/onnx/sentiment/").mkdir(parents=True, exist_ok=True)
         Path("models/onnx/sentiment_int8/").mkdir(parents=True, exist_ok=True)
         return
-        
+
     # 1. Aspect Extraction Model
     export_and_quantize(
         model_type="token_classification",
         source_dir=Path("models/aspect_extraction/best"),
         export_dir=Path("models/onnx/aspect_extraction"),
-        quantize_dir=Path("models/onnx/aspect_extraction_int8")
+        quantize_dir=Path("models/onnx/aspect_extraction_int8"),
     )
-    
+
     # 2. Sentiment Model (Multilingual)
     export_and_quantize(
         model_type="sequence_classification",
         source_dir=Path("models/sentiment/multilingual/best"),
         export_dir=Path("models/onnx/sentiment"),
-        quantize_dir=Path("models/onnx/sentiment_int8")
+        quantize_dir=Path("models/onnx/sentiment_int8"),
     )
 
+
 if __name__ == "__main__":
     main()

src/models/train_aspect_extraction.py

@@ -1,5 +1,3 @@
-import os
-import torch
 import numpy as np
 from pathlib import Path
 from datasets import load_from_disk
@@ -9,28 +7,28 @@ from transformers import (
     Trainer,
     DataCollatorForTokenClassification,
     AutoTokenizer,
-    set_seed
+    set_seed,
 )
 from seqeval.metrics import f1_score as seqeval_f1_score
-from seqeval.metrics import classification_report
 import mlflow
 
 from src.training.mlflow_utils import setup_mlflow
 
+
 def compute_metrics(p):
     """Computes evaluation metrics (F1 score) for token classification.
-    
+
     Args:
         p: EvalPrediction tuple containing predictions and labels.
-        
+
     Returns:
         Dictionary with 'f1' key and its computed value.
     """
     predictions, labels = p
     predictions = np.argmax(predictions, axis=2)
-    
+
     label_map = {0: "O", 1: "B-ASP", 2: "I-ASP"}
-    
+
     true_predictions = [
         [label_map[p] for (p, l) in zip(prediction, label) if l != -100]
         for prediction, label in zip(predictions, labels)
@@ -39,43 +37,42 @@ def compute_metrics(p):
         [label_map[l] for (p, l) in zip(prediction, label) if l != -100]
         for prediction, label in zip(predictions, labels)
     ]
-    
+
     # Seqeval F1 handles span-level scoring
     f1 = seqeval_f1_score(true_labels, true_predictions)
-    
+
     # Calculate macro F1 roughly from classification report if needed,
     # but for NER, seqeval's micro-averaged F1 (which seqeval_f1_score returns) is standard span-F1
-    return {
-        "f1": f1
-    }
+    return {"f1": f1}
+
 
 def main():
     """Main function to train and evaluate the aspect extraction model.
-    
+
     Loads tokenized dataset, initializes XLM-RoBERTa for token classification,
     configures Trainer, executes training loop, evaluates on test set,
     and logs results to MLflow.
     """
     set_seed(42)
     setup_mlflow()
-    
+
     dataset_path = Path("data/tokenized/absa_ner_dataset")
     print(f"Loading dataset from {dataset_path}")
     dataset = load_from_disk(str(dataset_path))
-    
+
     tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
     data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)
-    
+
     label_map = {0: "O", 1: "B-ASP", 2: "I-ASP"}
     model = AutoModelForTokenClassification.from_pretrained(
         "xlm-roberta-base",
         num_labels=len(label_map),
         id2label=label_map,
-        label2id={v: k for k, v in label_map.items()}
+        label2id={v: k for k, v in label_map.items()},
     )
-    
+
     output_dir = "models/aspect_extraction"
-    
+
     training_args = TrainingArguments(
         output_dir=output_dir,
         learning_rate=2e-5,
@@ -89,9 +86,9 @@ def main():
         metric_for_best_model="eval_f1",
         load_best_model_at_end=True,
         seed=42,
-        report_to="mlflow"
+        report_to="mlflow",
     )
-    
+
     trainer = Trainer(
         model=model,
         args=training_args,
@@ -99,28 +96,34 @@ def main():
         eval_dataset=dataset["validation"],
         tokenizer=tokenizer,
         data_collator=data_collator,
-        compute_metrics=compute_metrics
+        compute_metrics=compute_metrics,
     )
-    
+
     print("Training Aspect Extraction model...")
     trainer.train()
-    
+
     print("Evaluating on test set...")
     test_results = trainer.evaluate(dataset["test"], metric_key_prefix="test")
     print(test_results)
-    
+
     best_model_path = Path(output_dir) / "best"
     trainer.save_model(str(best_model_path))
     print(f"Best model saved to {best_model_path}")
-    
-    # Log test metric manually since trainer.train() only automatically logs eval metrics 
+
+    # 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.
-    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) as run:
-        mlflow.log_metrics({
-            "test_f1": test_results["test_f1"],
-            "test_loss": test_results["test_loss"]
-        })
+    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
+        )
+    ) as run:
+        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}")
 
+
 if __name__ == "__main__":
     main()

src/models/train_joint_absa.py

@@ -1,7 +1,7 @@
 """
 Script for training a Joint ABSA model (token classification + sentiment classification).
 """
-import os
+
 from pathlib import Path
 import torch
 import torch.nn as nn
@@ -11,41 +11,45 @@ from transformers import (
     AutoTokenizer,
     TrainingArguments,
     Trainer,
-    set_seed
+    set_seed,
 )
-from datasets import load_dataset
 import mlflow
 import numpy as np
 from sklearn.metrics import f1_score
-from transformers.modeling_outputs import TokenClassifierOutput, SequenceClassifierOutput
+from transformers.modeling_outputs import (
+    TokenClassifierOutput,
+    SequenceClassifierOutput,
+)
 from dataclasses import dataclass
 from typing import Optional, Tuple
 
 set_seed(42)
 
+
 @dataclass
 class JointModelOutput(TokenClassifierOutput, SequenceClassifierOutput):
     loss: Optional[torch.FloatTensor] = None
-    ner_logits: torch.FloatTensor = None
-    cls_logits: torch.FloatTensor = None
+    ner_logits: Optional[torch.FloatTensor] = None
+    cls_logits: Optional[torch.FloatTensor] = None
     hidden_states: Optional[Tuple[torch.FloatTensor]] = None
     attentions: Optional[Tuple[torch.FloatTensor]] = None
 
+
 class JointABSAModel(XLMRobertaPreTrainedModel):
     def __init__(self, config, num_ner_labels=3, num_sentiment_labels=4):
         super().__init__(config)
         self.num_ner_labels = num_ner_labels
         self.num_sentiment_labels = num_sentiment_labels
-        
+
         self.roberta = XLMRobertaModel(config, add_pooling_layer=False)
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        
+
         # Head 1: Token Classification (NER for Aspect Extraction)
         self.ner_classifier = nn.Linear(config.hidden_size, num_ner_labels)
-        
+
         # Head 2: Sequence Classification (Sentiment)
         self.sentiment_classifier = nn.Linear(config.hidden_size, num_sentiment_labels)
-        
+
         self.post_init()
 
     def forward(
@@ -56,13 +60,15 @@ class JointABSAModel(XLMRobertaPreTrainedModel):
         position_ids=None,
         head_mask=None,
         inputs_embeds=None,
-        labels=None, # NER labels
-        sentiment_labels=None, # Sentiment labels
+        labels=None,  # NER labels
+        sentiment_labels=None,  # Sentiment labels
         output_attentions=None,
         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,
@@ -78,29 +84,34 @@ class JointABSAModel(XLMRobertaPreTrainedModel):
 
         sequence_output = outputs[0]
         sequence_output = self.dropout(sequence_output)
-        
+
         # NER logits
         ner_logits = self.ner_classifier(sequence_output)
-        
+
         # Sentiment logits (using CLS token)
         cls_output = sequence_output[:, 0, :]
         cls_logits = self.sentiment_classifier(cls_output)
-        
+
         loss = None
         if labels is not None and sentiment_labels is not None:
             loss_fct = nn.CrossEntropyLoss()
-            
+
             # NER Loss
             active_loss = attention_mask.view(-1) == 1
             active_logits = ner_logits.view(-1, self.num_ner_labels)
             active_labels = torch.where(
-                active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+                active_loss,
+                labels.view(-1),
+                torch.tensor(loss_fct.ignore_index).type_as(labels),
             )
             ner_loss = loss_fct(active_logits, active_labels)
-            
+
             # Sentiment Loss
-            cls_loss = loss_fct(cls_logits.view(-1, self.num_sentiment_labels), sentiment_labels.view(-1))
-            
+            cls_loss = loss_fct(
+                cls_logits.view(-1, self.num_sentiment_labels),
+                sentiment_labels.view(-1),
+            )
+
             # Combined Loss
             loss = 0.5 * ner_loss + 0.5 * cls_loss
 
@@ -116,6 +127,7 @@ class JointABSAModel(XLMRobertaPreTrainedModel):
             attentions=outputs.attentions,
         )
 
+
 class JointTrainer(Trainer):
     def compute_loss(self, model, inputs, return_outputs=False):
         labels = inputs.pop("labels")
@@ -124,40 +136,45 @@ class JointTrainer(Trainer):
         loss = outputs.loss
         return (loss, outputs) if return_outputs else loss
 
+
 def compute_metrics(eval_pred) -> dict:
     # eval_pred.predictions is a tuple: (ner_logits, cls_logits)
     ner_logits, cls_logits = eval_pred.predictions
-    ner_labels = 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 
-    
+    ner_labels = 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
     # ... In a real setup, handle label pairing ...
-    
+
     cls_predictions = np.argmax(cls_logits, axis=-1)
     # Placeholder for joint span f1 logic
-    joint_span_f1 = 0.75 
-    
+    joint_span_f1 = 0.75
+
     # if sentiment_labels is available
     if sentiment_labels is not None:
         joint_macro_f1 = f1_score(sentiment_labels, cls_predictions, average="macro")
     else:
         joint_macro_f1 = 0.80
-        
-    return {
-        "joint_span_f1": joint_span_f1,
-        "joint_macro_f1": joint_macro_f1
-    }
+
+    return {"joint_span_f1": joint_span_f1, "joint_macro_f1": joint_macro_f1}
+
 
 def main():
     model_name = "xlm-roberta-base"
     output_dir = Path("models/joint_absa/best")
     output_dir.mkdir(parents=True, exist_ok=True)
-    
+
     print("Loading tokenizer and model...")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = JointABSAModel.from_pretrained(model_name, num_ner_labels=3, num_sentiment_labels=4)
-    
+    model = JointABSAModel.from_pretrained(
+        model_name, num_ner_labels=3, num_sentiment_labels=4
+    )
+
     training_args = TrainingArguments(
         output_dir=str(output_dir),
         evaluation_strategy="epoch",
@@ -167,24 +184,27 @@ def main():
         num_train_epochs=3,
         weight_decay=0.01,
         seed=42,
-        logging_dir='./logs',
+        logging_dir="./logs",
         logging_steps=10,
-        save_strategy="epoch"
+        save_strategy="epoch",
     )
-    
+
     # Placeholder dataset setup
     # In practice, need a DataCollator that handles both `labels` and `sentiment_labels`
-    
+
     mlflow.set_tracking_uri("sqlite:///mlflow.db")
     mlflow.set_experiment("joint-absa-training")
-    
+
     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)
         mlflow.log_metric("joint_macro_f1", 0.0)
 
+
 if __name__ == "__main__":
     main()

src/models/train_multilingual.py

@@ -1,9 +1,8 @@
 """
 Script for multilingual fine-tuning of XLM-RoBERTa using language-aware sampling.
 """
-import os
+
 from pathlib import Path
-import torch
 from torch.utils.data import WeightedRandomSampler
 from transformers import (
     AutoModelForSequenceClassification,
@@ -11,7 +10,7 @@ from transformers import (
     TrainingArguments,
     Trainer,
     DataCollatorWithPadding,
-    set_seed
+    set_seed,
 )
 from datasets import load_dataset, concatenate_datasets
 import mlflow
@@ -20,73 +19,83 @@ from sklearn.metrics import f1_score
 
 set_seed(42)
 
+
 def compute_metrics(eval_pred) -> dict:
     predictions, labels = eval_pred
     predictions = np.argmax(predictions, axis=1)
     return {"macro_f1": f1_score(labels, predictions, average="macro")}
 
+
 class LanguageAwareTrainer(Trainer):
     def _get_train_sampler(self):
         dataset = self.train_dataset
-        
+
         # Calculate weights to achieve 1:1 English:Hindi ratio
         # Assuming dataset has a 'lang' feature
-        lang_labels = dataset['lang']
-        en_count = sum(1 for l in lang_labels if l == 'en')
-        hi_count = sum(1 for l in lang_labels if l == 'hi')
-        
+        lang_labels = dataset["lang"]
+        en_count = sum(1 for l in lang_labels if l == "en")
+        hi_count = sum(1 for l in lang_labels if l == "hi")
+
         weights = []
         for l in lang_labels:
-            if l == 'en':
+            if l == "en":
                 weights.append(1.0 / en_count if en_count > 0 else 0)
-            elif l == 'hi':
+            elif l == "hi":
                 weights.append(1.0 / hi_count if hi_count > 0 else 0)
             else:
                 weights.append(0)
-                
+
         # WeightedRandomSampler handles the sampling
-        return WeightedRandomSampler(weights, num_samples=len(dataset), replacement=True)
+        return WeightedRandomSampler(
+            weights, num_samples=len(dataset), replacement=True
+        )
+
 
 def main():
     model_name = "xlm-roberta-base"
     output_dir = Path("models/sentiment/multilingual/best")
     output_dir.mkdir(parents=True, exist_ok=True)
-    
+
     data_dir = Path("data/processed")
     en_train_file = data_dir / "semeval_train.jsonl"
     hi_train_file = data_dir / "hindi_augmented.jsonl"
-    
+
     # NOTE: Dummy loading handling for execution without actual files
     if not en_train_file.exists() or not hi_train_file.exists():
-        print("Missing dataset files. Ensure SemEval and Hindi augmented files are present.")
+        print(
+            "Missing dataset files. Ensure SemEval and Hindi augmented files are present."
+        )
         return
 
     print("Loading datasets...")
     en_dataset = load_dataset("json", data_files={"train": str(en_train_file)})["train"]
     hi_dataset = load_dataset("json", data_files={"train": str(hi_train_file)})["train"]
-    
+
     # Ensure they have a 'lang' column for our sampler
     def add_en_lang(example):
-        example['lang'] = 'en'
+        example["lang"] = "en"
         return example
+
     def add_hi_lang(example):
-        example['lang'] = 'hi'
+        example["lang"] = "hi"
         return example
-        
+
     en_dataset = en_dataset.map(add_en_lang)
     hi_dataset = hi_dataset.map(add_hi_lang)
-    
+
     train_dataset = concatenate_datasets([en_dataset, hi_dataset])
-    
+
     print("Loading tokenizer and model...")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=4)
-    
+
     def tokenize_function(examples):
-        return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
-        
+        return tokenizer(
+            examples["text"], truncation=True, padding="max_length", max_length=128
+        )
+
     tokenized_train = train_dataset.map(tokenize_function, batched=True)
-    
+
     training_args = TrainingArguments(
         output_dir=str(output_dir),
         evaluation_strategy="epoch",
@@ -96,39 +105,40 @@ def main():
         num_train_epochs=3,
         weight_decay=0.01,
         seed=42,
-        save_strategy="epoch"
+        save_strategy="epoch",
     )
-    
+
     trainer = LanguageAwareTrainer(
         model=model,
         args=training_args,
         train_dataset=tokenized_train,
         tokenizer=tokenizer,
         data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
-        compute_metrics=compute_metrics
+        compute_metrics=compute_metrics,
     )
-    
+
     mlflow.set_tracking_uri("sqlite:///mlflow.db")
     mlflow.set_experiment("multilingual-sentiment")
-    
+
     with mlflow.start_run():
         print("Starting multilingual training...")
         # trainer.train() # Uncomment to run actual training
-        
+
         # NOTE: Placeholder for evaluation logging
         en_f1 = 0.82
         hi_f1 = 0.68
         combined_f1 = 0.75
         gap = en_f1 - hi_f1
-        
+
         mlflow.log_metric("en_macro_f1", en_f1)
         mlflow.log_metric("hi_macro_f1", hi_f1)
         mlflow.log_metric("combined_macro_f1", combined_f1)
         mlflow.log_metric("cross_lingual_gap", gap)
-        
+
         print("Saving model...")
         model.save_pretrained(str(output_dir))
         tokenizer.save_pretrained(str(output_dir))
 
+
 if __name__ == "__main__":
     main()

src/models/train_qlora.py

@@ -1,7 +1,7 @@
 """
 Script for QLoRA fine-tuning of XLM-RoBERTa for sentiment analysis.
 """
-import os
+
 from pathlib import Path
 import torch
 from transformers import (
@@ -11,7 +11,7 @@ from transformers import (
     TrainingArguments,
     Trainer,
     DataCollatorWithPadding,
-    set_seed
+    set_seed,
 )
 from peft import get_peft_model, LoraConfig, TaskType
 from datasets import load_dataset
@@ -22,38 +22,42 @@ from sklearn.metrics import f1_score
 # Constraints: seed=42 everywhere
 set_seed(42)
 
+
 def compute_metrics(eval_pred) -> dict:
     predictions, labels = eval_pred
     predictions = np.argmax(predictions, axis=1)
     macro_f1 = f1_score(labels, predictions, average="macro")
     return {"macro_f1": macro_f1}
 
+
 def main():
     model_name = "xlm-roberta-base"
     output_dir = Path("models/sentiment/qlora-adapter")
     output_dir.mkdir(parents=True, exist_ok=True)
-    
+
     data_dir = Path("data/processed")
-    
+
     # 4-bit quantization config
     try:
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_compute_dtype=torch.float16,
             bnb_4bit_quant_type="nf4",
-            bnb_4bit_use_double_quant=True
+            bnb_4bit_use_double_quant=True,
         )
     except Exception as e:
-        print(f"Warning: bitsandbytes might not be supported on this system. Detailed error: {e}")
-        bnb_config = None # Fallback or error based on environment
+        print(
+            f"Warning: bitsandbytes might not be supported on this system. Detailed error: {e}"
+        )
+        bnb_config = None  # Fallback or error based on environment
 
     print("Loading tokenizer and model...")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(
         model_name,
-        num_labels=4, # positive, negative, neutral, conflict
+        num_labels=4,  # positive, negative, neutral, conflict
         quantization_config=bnb_config if bnb_config else None,
-        device_map="auto"
+        device_map="auto",
     )
 
     lora_config = LoraConfig(
@@ -61,12 +65,12 @@ def main():
         r=16,
         lora_alpha=32,
         lora_dropout=0.1,
-        target_modules=["query", "value"]
+        target_modules=["query", "value"],
     )
-    
+
     model = get_peft_model(model, lora_config)
     model.print_trainable_parameters()
-    
+
     # NOTE: Assuming combined dataset is prepared or we combine them here.
     # For now, we load a placeholder train dataset
     train_file = data_dir / "semeval_train.jsonl"
@@ -75,12 +79,14 @@ def main():
         return
 
     dataset = load_dataset("json", data_files={"train": str(train_file)})
-    
+
     def tokenize_function(examples):
-        return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
-        
+        return tokenizer(
+            examples["text"], truncation=True, padding="max_length", max_length=128
+        )
+
     tokenized_datasets = dataset.map(tokenize_function, batched=True)
-    
+
     training_args = TrainingArguments(
         output_dir=str(output_dir),
         evaluation_strategy="epoch",
@@ -90,11 +96,11 @@ def main():
         num_train_epochs=3,
         weight_decay=0.01,
         seed=42,
-        logging_dir='./logs',
+        logging_dir="./logs",
         logging_steps=10,
-        save_strategy="epoch"
+        save_strategy="epoch",
     )
-    
+
     trainer = Trainer(
         model=model,
         args=training_args,
@@ -102,21 +108,22 @@ def main():
         # eval_dataset=tokenized_datasets["test"], # Add test set if available
         tokenizer=tokenizer,
         data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
-        compute_metrics=compute_metrics
+        compute_metrics=compute_metrics,
     )
-    
+
     mlflow.set_tracking_uri("sqlite:///mlflow.db")
     mlflow.set_experiment("qlora-sentiment")
-    
+
     with mlflow.start_run():
         trainer.train()
-        
+
         # Save adapter
         model.save_pretrained(str(output_dir))
         tokenizer.save_pretrained(str(output_dir))
-        
+
         # Log adapter weights to MLflow
         mlflow.log_artifacts(str(output_dir), artifact_path="qlora-adapter")
 
+
 if __name__ == "__main__":
     main()

src/models/train_sentiment.py

@@ -1,4 +1,3 @@
-import os
 import torch
 import numpy as np
 from pathlib import Path
@@ -9,28 +8,29 @@ from transformers import (
     Trainer,
     DataCollatorWithPadding,
     AutoTokenizer,
-    set_seed
+    set_seed,
 )
 from sklearn.metrics import f1_score, confusion_matrix
 import mlflow
 
 from src.training.mlflow_utils import setup_mlflow
 
+
 def compute_metrics(p):
     """Computes evaluation metrics (F1 score) for sequence classification.
-    
+
     Args:
         p: EvalPrediction tuple containing predictions and labels.
-        
+
     Returns:
         Dictionary with macro F1 and per-class F1 metrics.
     """
     predictions, labels = p
     predictions = np.argmax(predictions, axis=1)
-    
+
     macro_f1 = f1_score(labels, predictions, average="macro")
     per_class_f1 = f1_score(labels, predictions, average=None)
-    
+
     # We will log confusion matrix in the main function
     return {
         "macro_f1": macro_f1,
@@ -40,52 +40,56 @@ def compute_metrics(p):
         "f1_conflict": per_class_f1[3] if len(per_class_f1) > 3 else 0.0,
     }
 
+
 class ImbalancedTrainer(Trainer):
     def __init__(self, class_weights=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.class_weights = class_weights
-        
+
     def compute_loss(self, model, inputs, return_outputs=False):
         labels = inputs.pop("labels")
         outputs = model(**inputs)
         logits = outputs.logits
-        
+
         if self.class_weights is not None:
-            loss_fct = torch.nn.CrossEntropyLoss(weight=self.class_weights.to(model.device))
+            loss_fct = torch.nn.CrossEntropyLoss(
+                weight=self.class_weights.to(model.device)
+            )
         else:
             loss_fct = torch.nn.CrossEntropyLoss()
-            
+
         loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
-        
+
         return (loss, outputs) if return_outputs else loss
 
+
 def main():
     """Main function to train and evaluate the sentiment classification model.
-    
+
     Loads tokenized dataset, initializes XLM-RoBERTa for sequence classification,
     handles class imbalances using a custom Trainer, executes training loop,
     evaluates on test set, logs confusion matrix, and logs results to MLflow.
     """
     set_seed(42)
     setup_mlflow()
-    
+
     dataset_path = Path("data/tokenized/absa_cls_dataset")
     print(f"Loading dataset from {dataset_path}")
     dataset = load_from_disk(str(dataset_path))
-    
+
     tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
-    
+
     label_map = {0: "positive", 1: "negative", 2: "neutral", 3: "conflict"}
     model = AutoModelForSequenceClassification.from_pretrained(
         "xlm-roberta-base",
         num_labels=len(label_map),
         id2label=label_map,
-        label2id={v: k for k, v in label_map.items()}
+        label2id={v: k for k, v in label_map.items()},
     )
-    
+
     output_dir = "models/sentiment"
-    
+
     training_args = TrainingArguments(
         output_dir=output_dir,
         learning_rate=2e-5,
@@ -99,15 +103,18 @@ def main():
         metric_for_best_model="eval_macro_f1",
         load_best_model_at_end=True,
         seed=42,
-        report_to="mlflow"
+        report_to="mlflow",
     )
-    
+
     # Calculate class weights for imbalanced dataset (especially 'conflict')
     train_labels = dataset["train"]["label"]
     from sklearn.utils.class_weight import compute_class_weight
-    class_weights = compute_class_weight('balanced', classes=np.unique(train_labels), y=train_labels)
+
+    class_weights = compute_class_weight(
+        "balanced", classes=np.unique(train_labels), y=train_labels
+    )
     class_weights_tensor = torch.tensor(class_weights, dtype=torch.float)
-    
+
     trainer = ImbalancedTrainer(
         model=model,
         args=training_args,
@@ -116,33 +123,42 @@ def main():
         tokenizer=tokenizer,
         data_collator=data_collator,
         compute_metrics=compute_metrics,
-        class_weights=class_weights_tensor
+        class_weights=class_weights_tensor,
     )
-    
+
     print("Training Sentiment Classification model...")
     trainer.train()
-    
+
     print("Evaluating on test set...")
     test_results = trainer.evaluate(dataset["test"], metric_key_prefix="test")
     print(test_results)
-    
+
     best_model_path = Path(output_dir) / "best"
     trainer.save_model(str(best_model_path))
     print(f"Best model saved to {best_model_path}")
-    
+
     # Confusion matrix on test set
     predictions = trainer.predict(dataset["test"])
     preds = np.argmax(predictions.predictions, axis=1)
     labels = predictions.label_ids
     cm = confusion_matrix(labels, preds)
-    
-    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) as run:
-        mlflow.log_metrics({
-            "test_macro_f1": test_results["test_macro_f1"],
-            "test_loss": test_results["test_loss"]
-        })
+
+    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
+        )
+    ) as run:
+        mlflow.log_metrics(
+            {
+                "test_macro_f1": test_results["test_macro_f1"],
+                "test_loss": test_results["test_loss"],
+            }
+        )
         mlflow.log_dict({"confusion_matrix": cm.tolist()}, "confusion_matrix.json")
         print(f"Logged test metrics and confusion matrix to run {run.info.run_id}")
 
+
 if __name__ == "__main__":
     main()

src/training/mlflow_utils.py

@@ -1,73 +1,84 @@
 import mlflow
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Union
 from pathlib import Path
-import os
 
 # Default configuration
 MLFLOW_TRACKING_URI = "sqlite:///mlflow/mlflow.db"
 EXPERIMENT_NAME = "multilingual-absa"
 
+
 def setup_mlflow():
     """Initializes MLflow tracking URI and experiment."""
     # Ensure the directory exists
     Path("mlflow").mkdir(parents=True, exist_ok=True)
-    
+
     mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
     mlflow.set_experiment(EXPERIMENT_NAME)
 
-def log_training_run(params: Dict[str, Any], metrics: Dict[str, float], model_path: Optional[str | Path] = None, run_name: Optional[str] = None) -> str:
+
+def log_training_run(
+    params: Dict[str, Any],
+    metrics: Dict[str, float],
+    model_path: Optional[Union[str, Path]] = None,
+    run_name: Optional[str] = None,
+) -> str:
     """
     Logs parameters, metrics, and optionally a model artifact to MLflow.
-    
+
     Args:
         params: Dictionary of hyperparameters or configuration.
         metrics: Dictionary of evaluation metrics.
         model_path: Path to the saved model directory or file.
         run_name: Optional name for the run.
-        
+
     Returns:
         The ID of the created MLflow run.
     """
     setup_mlflow()
-    
+
     with mlflow.start_run(run_name=run_name) as run:
         mlflow.log_params(params)
         mlflow.log_metrics(metrics)
-        
+
         if model_path:
             model_path_obj = Path(model_path)
             if model_path_obj.exists():
                 mlflow.log_artifact(str(model_path_obj), artifact_path="model")
             else:
-                print(f"Warning: Model path {model_path} does not exist. Artifact not logged.")
-                
+                print(
+                    f"Warning: Model path {model_path} does not exist. Artifact not logged."
+                )
+
         return run.info.run_id
 
-def get_best_run(metric: str = "eval_macro_f1", ascending: bool = False) -> Optional[mlflow.entities.Run]:
+
+def get_best_run(
+    metric: str = "eval_macro_f1", ascending: bool = False
+) -> Optional[Any]: # type: ignore
     """
     Retrieves the best run from the experiment based on a specific metric.
-    
+
     Args:
         metric: The metric to sort by.
         ascending: True if a lower metric is better (e.g., loss), False for higher is better (e.g., F1).
-        
+
     Returns:
         The MLflow Run object for the best run, or None if no runs exist.
     """
     setup_mlflow()
-    
+
     experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
     if not experiment:
         return None
-        
+
     runs = mlflow.search_runs(
         experiment_ids=[experiment.experiment_id],
         order_by=[f"metrics.{metric} {'ASC' if ascending else 'DESC'}"],
         max_results=1,
-        output_format="list"
+        output_format="list",
     )
-    
+
     if not runs:
         return None
-        
+
     return runs[0]