Added local vs HF mode

.gitignore

@@ -1,6 +1,6 @@
 ProjectEnv
-saved_model
+models
 .pytest_cache
-artifacts
 data/__pycache__
-data/dataset
+data/dataset
+app/__pycache__

src/app/app.py

@@ -9,7 +9,7 @@ import csv
 import requests
 from typing import Union, List
 import torch
-
+from .config import MODEL_SOURCE, ModelSource, load_model_and_tokenizer
 
 
 app = FastAPI()
@@ -18,17 +18,13 @@ app = FastAPI()
 class SentimentQuery(BaseModel):
     input_texts: Union[str, List[str]]
 
-task='sentiment'
-mapping_link = f"https://raw.githubusercontent.com/cardiffnlp/tweeteval/main/datasets/{task}/mapping.txt"
+mapping_link = f"https://raw.githubusercontent.com/cardiffnlp/tweeteval/main/datasets/sentiment/mapping.txt"
 with urllib.request.urlopen(mapping_link) as f:
     html = f.read().decode('utf-8').split("\n")
     csvreader = csv.reader(html, delimiter='\t')
 labels = [row[1] for row in csvreader if len(row) > 1]
 
-MODEL = f"cardiffnlp/twitter-roberta-base-{task}-latest"
-model = AutoModelForSequenceClassification.from_pretrained(MODEL)
-tokenizer = AutoTokenizer.from_pretrained(MODEL)
-
+tokenizer, model = load_model_and_tokenizer(MODEL_SOURCE)
 
 
 @app.post("/predict")

src/app/config.py

@@ -0,0 +1,24 @@
+import os
+from enum import Enum
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from pathlib import Path
+
+
+class ModelSource(str, Enum):
+    HF = "hf"
+    LOCAL = "local"
+
+MODEL_SOURCE = ModelSource(os.getenv("MODEL_SOURCE", "hf"))
+HF_MODEL = f"cardiffnlp/twitter-roberta-base-sentiment-latest"
+
+
+def load_model_and_tokenizer(MODEL_SOURCE):
+    if MODEL_SOURCE == ModelSource.HF:   # use the latest model available in the HF hub
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL)
+        model = AutoModelForSequenceClassification.from_pretrained(HF_MODEL)
+    else: # use a locally fine tuned model
+        local_model_path = Path("models/saved_model")
+        assert local_model_path.exists(), """No local model was found. Run 'python3 src/train_model.py'"""
+        tokenizer = AutoTokenizer.from_pretrained("models/saved_tokenizer")
+        model = AutoModelForSequenceClassification.from_pretrained("models/saved_model")
+    return tokenizer, model

src/train_model.py

@@ -0,0 +1,241 @@
+from app.utils import preprocess
+import urllib
+import csv
+import os
+import torch
+from transformers import (
+    AutoTokenizer, AutoModelForSequenceClassification,
+    TrainingArguments, Trainer, EarlyStoppingCallback,
+    DataCollatorWithPadding
+)
+from datasets import load_from_disk
+
+
+# --- Device detection ---
+if torch.cuda.is_available():
+    device = "cuda"
+    use_bf16 = torch.cuda.is_bf16_supported()
+    use_fp16 = not use_bf16
+elif torch.backends.mps.is_available():
+    device = "mps"
+    use_bf16 = False
+    use_fp16 = False
+else:
+    device = "cpu"
+    use_bf16 = False
+    use_fp16 = False
+    
+if device == "cuda" and use_bf16:
+    load_dtype = torch.bfloat16
+elif device == "cuda" and use_fp16:
+    load_dtype = torch.float16
+else:
+    load_dtype = torch.float32  # MPS/CPU -> fp32
+
+import evaluate
+
+
+MODEL = f"cardiffnlp/twitter-roberta-base-sentiment"
+
+# download label mapping
+mapping_link = f"https://raw.githubusercontent.com/cardiffnlp/tweeteval/main/datasets/sentiment/mapping.txt"
+with urllib.request.urlopen(mapping_link) as f:
+    html = f.read().decode('utf-8').split("\n")
+    csvreader = csv.reader(html, delimiter='\t')
+labels = [row[1] for row in csvreader if len(row) > 1]
+
+
+
+# --- Tokenizer: keep short max_length to save memory ---
+tokenizer = AutoTokenizer.from_pretrained(MODEL, use_fast=True, model_max_length=128)
+
+
+def tokenize_function(batch):
+    return tokenizer(
+        batch["text"],
+        truncation=True,
+        max_length=128,
+        padding=False  # we will pad per-batch via DataCollatorWithPadding
+    )
+
+data_collator = DataCollatorWithPadding(
+    tokenizer=tokenizer,
+    pad_to_multiple_of=8 if (device == "cuda" and (use_bf16 or use_fp16)) else None
+)
+
+
+model = AutoModelForSequenceClassification.from_pretrained(
+    MODEL, num_labels=3, torch_dtype=load_dtype
+)
+
+model.gradient_checkpointing_enable()
+model.config.use_cache = False
+
+#### DATASET LOADING
+
+
+dataset_path = "data/dataset"  # same path you used before
+dataset = load_from_disk(dataset_path)
+
+
+# ---- COPY-PASTE FROM HERE ----
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+from datasets import DatasetDict
+from transformers import AutoTokenizer, DataCollatorWithPadding
+
+def make_trainer_ready(
+    raw_ds: DatasetDict,
+    model_name: str = "cardiffnlp/twitter-roberta-base-sep2022",
+    train_frac: float = 0.2,
+    val_frac: float = 0.2,
+    seed: int = 42,
+    label_col: str = "label",
+    text_col: str = "text",
+    max_length: int = 128,
+    pad_to_multiple_of_8_on_cuda: bool = True,
+):
+    """
+    Returns (train_ds, eval_ds, data_collator, tokenizer) ready for HF Trainer.
+    - Ensures there's a validation split (creates one from train if missing).
+    - Takes fractional subsets, stratified by label when possible.
+    - Tokenizes and keeps only the columns Trainer expects.
+    """
+    assert 0 < train_frac <= 1.0, "train_frac must be in (0,1]."
+    assert 0 < val_frac <= 1.0, "val_frac must be in (0,1]."
+    assert text_col in raw_ds["train"].column_names, f"Missing text column: {text_col}"
+    assert label_col in raw_ds["train"].column_names, f"Missing label column: {label_col}"
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, model_max_length=max_length)
+
+    # 1) Ensure we have a validation split
+    if "validation" not in raw_ds:
+        split = raw_ds["train"].train_test_split(
+            test_size=val_frac,
+            stratify_by_column=label_col if label_col in raw_ds["train"].column_names else None,
+            seed=seed,
+        )
+        raw_ds = DatasetDict(train=split["train"], validation=split["test"])
+    else:
+        raw_ds = DatasetDict(train=raw_ds["train"], validation=raw_ds["validation"])
+
+    # 2) Take fractions (stratified when possible)
+    def take_frac(ds, frac):
+        if frac >= 1.0:  # keep full split
+            return ds
+        out = ds.train_test_split(
+            test_size=1 - frac,
+            stratify_by_column=label_col if label_col in ds.column_names else None,
+            seed=seed,
+        )
+        return out["train"]  # the kept fraction
+
+    small_train = take_frac(raw_ds["train"], train_frac)
+    small_eval  = take_frac(raw_ds["validation"], val_frac)
+
+    # 3) Tokenize (no padding here; we pad per-batch with the collator)
+    def tok(batch):
+        return tokenizer(batch[text_col], truncation=True, max_length=max_length, padding=False)
+
+    small_train_tok = small_train.map(tok, batched=True, remove_columns=[c for c in small_train.column_names if c not in (text_col, label_col)])
+    small_eval_tok  = small_eval.map(tok,  batched=True, remove_columns=[c for c in small_eval.column_names  if c not in (text_col, label_col)])
+
+    # 4) Keep only the columns Trainer needs
+    keep_cols = ["input_ids", "attention_mask", label_col]
+    small_train_tok = small_train_tok.remove_columns([c for c in small_train_tok.column_names if c not in keep_cols])
+    small_eval_tok  = small_eval_tok.remove_columns([c for c in small_eval_tok.column_names  if c not in keep_cols])
+
+    # 5) Data collator with dynamic padding (CUDA gets pad_to_multiple_of=8)
+    import torch
+    pad_to_mult = 8 if (pad_to_multiple_of_8_on_cuda and torch.cuda.is_available()) else None
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer, pad_to_multiple_of=pad_to_mult)
+
+    return small_train_tok, small_eval_tok, data_collator, tokenizer
+
+
+train_ds, eval_ds, data_collator, tokenizer = make_trainer_ready(
+    raw_ds=dataset,
+    model_name="cardiffnlp/twitter-roberta-base-sep2022",
+    train_frac=0.2,    # take 20% of train
+    val_frac=0.5,      # take 50% of validation
+    seed=42,
+    label_col="label",
+    text_col="text",
+    max_length=128,
+)
+
+# --- Training args: stop forking on macOS, fix pin_memory ---
+trainer_fp16 = bool(device == "cuda" and use_fp16)
+trainer_bf16 = bool(device == "cuda" and use_bf16)
+
+training_args = TrainingArguments(
+    output_dir="models/artifacts",
+    learning_rate=1e-5,
+    per_device_train_batch_size=4,
+    per_device_eval_batch_size=8,
+    gradient_accumulation_steps=8,
+    num_train_epochs=3,
+    weight_decay=0.01,
+    warmup_ratio=0.1,
+    lr_scheduler_type="linear",
+
+    eval_strategy="steps",
+    logging_strategy="steps",
+    save_strategy="steps",
+    eval_steps=500,
+    logging_steps=100,
+    save_steps=500,
+
+    load_best_model_at_end=True,
+    metric_for_best_model="recall",
+    greater_is_better=True,
+    save_total_limit=2,
+
+    # Precision
+    fp16=trainer_fp16,
+    bf16=trainer_bf16,
+
+    # DataLoader knobs (avoid fork/tokenizers warning on macOS)
+    dataloader_num_workers=0,                         # <- key for macOS/MPS
+    dataloader_pin_memory=(device == "cuda"),         # False on MPS/CPU, True on CUDA
+    group_by_length=True,
+    report_to="none",
+)
+
+# --- Metrics (macro recall, etc.) ---
+recall_metric = evaluate.load("recall")
+acc_metric = evaluate.load("accuracy")
+f1_metric = evaluate.load("f1")
+
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    preds = logits.argmax(axis=-1)
+    return {
+        "accuracy": acc_metric.compute(predictions=preds, references=labels)["accuracy"],
+        "f1_macro": f1_metric.compute(predictions=preds, references=labels, average="macro")["f1"],
+        "recall": recall_metric.compute(predictions=preds, references=labels, average="macro")["recall"],
+    }
+
+callbacks = [EarlyStoppingCallback(early_stopping_patience=2)]
+
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset= train_ds,
+    eval_dataset= eval_ds,
+    compute_metrics=compute_metrics,
+    data_collator=data_collator,       # <- important
+    tokenizer=tokenizer,
+    callbacks=callbacks,
+)
+
+model.to(device)
+trainer.train()
+trainer.save_model("models/saved_model")
+tokenizer.save_pretrained("models/saved_tokenizer")
+try:
+    trainer.create_model_card()
+except Exception:
+    pass