Created src/dataset.py — converts Debate objects into PyTorch tensors for RoBERTa

Created src/model.py — loads roberta-base configured for 4-class classification
Created src/train.py — full training loop with class weights, AdamW optimizer, validation, checkpoint saving
Created src/__init__.py — predict() and predict_debate() functions for other team members to use
Fixed Person 1's IBM loader bug (was ignoring stance_WA, labeling everything as premise)
Trained a first model (wrong labels, now discarded)
What still needs to be done right now:

Retrain with the fixed labels — run python -m src.train --epochs 3 --batch-size 8
Sanity check the model once training finishes — run a quick predict test to confirm it returns different labels for different inputs
Commit and push the fixed code to the repo so the rest of the team can use predict() and predict_debate()

data/loaders/ibm.py

@@ -6,7 +6,9 @@ HuggingFace: ibm/argument_quality_ranking_30k
 
 We map this to the unified schema:
   - topic → Debate with a single claim (the motion)
-  - argument → Argument of type 'premise' attached to that claim
+  - stance_WA = 1  → premise      (argument supports the topic)
+  - stance_WA = -1 → counter_claim (argument opposes the topic)
+  - WA < 0.3       → unknown      (very low quality / unclear)
   - WA (weighted average quality) stored in metadata
 """
 
@@ -14,6 +16,14 @@ from typing import List, Dict
 
 from data.schema import Argument, Debate
 
+_UNKNOWN_QUALITY_THRESHOLD = 0.3
+
+
+def _arg_type(stance_wa: int, quality: float) -> str:
+    if quality < _UNKNOWN_QUALITY_THRESHOLD:
+        return "unknown"
+    return "premise" if stance_wa == 1 else "counter_claim"
+
 
 def load_ibm(split: str = "train") -> List[Debate]:
     try:
@@ -40,15 +50,15 @@ def load_ibm(split: str = "train") -> List[Debate]:
             text=topic,
             arg_type="claim",
         )
-        premises = [
+        arguments = [
             Argument(
                 id=f"{topic_id}__{i}",
                 text=row["argument"],
-                arg_type="premise",
+                arg_type=_arg_type(row["stance_WA"], row["WA"]),
                 parent_id=claim.id,
                 metadata={
-                    "quality_score": row.get("WA"),
-                    "stance": row.get("stance"),
+                    "quality_score": row["WA"],
+                    "stance_wa": row["stance_WA"],
                 },
             )
             for i, row in enumerate(rows)
@@ -58,7 +68,7 @@ def load_ibm(split: str = "train") -> List[Debate]:
                 id=topic_id,
                 title=topic,
                 source="ibm",
-                arguments=[claim] + premises,
+                arguments=[claim] + arguments,
             )
         )
     return debates

src/__init__.py

@@ -1,43 +1,62 @@
 from __future__ import annotations
 
-import torch
-from transformers import RobertaForSequenceClassification, RobertaTokenizerFast
-
 from data.schema import Argument, Debate
-from .dataset import ID2LABEL, LABEL2ID
+
+LABEL2ID = {"claim": 0, "counter_claim": 1, "premise": 2, "unknown": 3}
+ID2LABEL = {v: k for k, v in LABEL2ID.items()}
 
 _model = None
 _tokenizer = None
 _loaded_checkpoint: str | None = None
-_device = "cuda" if torch.cuda.is_available() else "cpu"
+_device: str | None = None
 
 
 def _load(checkpoint_dir: str) -> None:
-    global _model, _tokenizer, _loaded_checkpoint
+    global _model, _tokenizer, _loaded_checkpoint, _device
     if _loaded_checkpoint != checkpoint_dir:
+        import torch
+        from transformers import (
+            RobertaForSequenceClassification,
+            RobertaTokenizerFast,
+        )
+        _device = "cuda" if torch.cuda.is_available() else "cpu"
         _tokenizer = RobertaTokenizerFast.from_pretrained(checkpoint_dir)
-        _model = RobertaForSequenceClassification.from_pretrained(checkpoint_dir)
+        _model = RobertaForSequenceClassification.from_pretrained(
+            checkpoint_dir
+        )
         _model.eval()
         _model.to(_device)
         _loaded_checkpoint = checkpoint_dir
 
 
-def predict(text: str, parent_text: str = "", checkpoint_dir: str = "models/best") -> str:
+def predict(
+    text: str,
+    parent_text: str = "",
+    checkpoint_dir: str = "models/best",
+) -> str:
     """Classify a single argument text.
 
     Returns one of: 'claim', 'counter_claim', 'premise', 'unknown'.
-    Pass parent_text when the comment is a reply — the model uses both together.
+    Pass parent_text when the comment is a reply.
     """
+    import torch
     _load(checkpoint_dir)
     if parent_text:
         enc = _tokenizer(
-            parent_text, text,
-            return_tensors="pt", truncation=True, max_length=256, padding="max_length",
+            parent_text,
+            text,
+            return_tensors="pt",
+            truncation=True,
+            max_length=256,
+            padding="max_length",
         )
     else:
         enc = _tokenizer(
             text,
-            return_tensors="pt", truncation=True, max_length=256, padding="max_length",
+            return_tensors="pt",
+            truncation=True,
+            max_length=256,
+            padding="max_length",
         )
     enc = {k: v.to(_device) for k, v in enc.items()}
     with torch.no_grad():
@@ -45,12 +64,14 @@ def predict(text: str, parent_text: str = "", checkpoint_dir: str = "models/best
     return ID2LABEL[logits.argmax(dim=-1).item()]
 
 
-def predict_debate(debate: Debate, checkpoint_dir: str = "models/best") -> Debate:
-    """Classify every argument in a debate and return a new Debate with predicted labels.
+def predict_debate(
+    debate: Debate,
+    checkpoint_dir: str = "models/best",
+) -> Debate:
+    """Label every argument in a debate, return new Debate with predictions.
 
-    This is the main entry point for Person 3 (eval) and Person 4 (failure analysis).
-    The returned Debate has the same structure and parent_id links — just with
-    arg_type replaced by the model's predictions.
+    Main entry point for Person 3 (eval) and Person 4 (failure analysis).
+    Preserves structure and parent_id links; only arg_type is replaced.
     """
     _load(checkpoint_dir)
     arg_map = {a.id: a for a in debate.arguments}
@@ -60,7 +81,11 @@ def predict_debate(debate: Debate, checkpoint_dir: str = "models/best") -> Debat
         labeled.append(Argument(
             id=arg.id,
             text=arg.text,
-            arg_type=predict(arg.text, parent.text if parent else "", checkpoint_dir),
+            arg_type=predict(
+                arg.text,
+                parent.text if parent else "",
+                checkpoint_dir,
+            ),
             parent_id=arg.parent_id,
             author=arg.author,
             score=arg.score,

src/train.py

@@ -6,41 +6,38 @@ import os
 from collections import Counter
 from pathlib import Path
 
-import torch
-from torch.utils.data import DataLoader, random_split
-from torch.optim import AdamW
-from transformers import RobertaTokenizerFast, get_linear_schedule_with_warmup
-from sklearn.metrics import f1_score, accuracy_score
-from tqdm import tqdm
-
 from data import load_ibm, clean_debates
-from .dataset import ArgumentDataset, NUM_LABELS
-from .model import MODEL_NAME, build_model
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+NUM_LABELS = 4
 
 
-def _class_weights(dataset: ArgumentDataset) -> torch.Tensor:
-    """Inverse-frequency weights so rare labels (counter_claim, unknown) aren't ignored."""
+def _class_weights(dataset, device):
+    """Inverse-frequency weights so rare labels aren't ignored."""
+    import torch
     counts = Counter(ex["label"] for ex in dataset.examples)
     total = sum(counts.values())
     weights = torch.ones(NUM_LABELS)
     for label_id, count in counts.items():
         weights[label_id] = total / (NUM_LABELS * count)
-    return weights
+    return weights.to(device)
 
 
-def evaluate(model, loader) -> tuple[float, float, float]:
+def evaluate(model, loader, device):
+    import torch
+    from sklearn.metrics import f1_score, accuracy_score
     model.eval()
     all_preds, all_labels = [], []
     total_loss = 0.0
     loss_fn = torch.nn.CrossEntropyLoss()
     with torch.no_grad():
         for batch in loader:
-            input_ids = batch["input_ids"].to(DEVICE)
-            attention_mask = batch["attention_mask"].to(DEVICE)
-            labels = batch["label"].to(DEVICE)
-            logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
+            input_ids = batch["input_ids"].to(device)
+            attention_mask = batch["attention_mask"].to(device)
+            labels = batch["label"].to(device)
+            logits = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            ).logits
             total_loss += loss_fn(logits, labels).item()
             all_preds.extend(logits.argmax(dim=-1).cpu().numpy())
             all_labels.extend(labels.cpu().numpy())
@@ -61,7 +58,7 @@ def train(
 ) -> str:
     Path(output_dir).mkdir(parents=True, exist_ok=True)
 
-    print(f"Device: {DEVICE}")
+    # --- load data first, before touching CUDA ---
     print("Loading data...")
     debates = clean_debates(load_ibm("train"))
     print(f"  IBM Debater: {len(debates)} debates")
@@ -70,9 +67,37 @@ def train(
         from data import load_cmv
         cmv = clean_debates(load_cmv("train"))
         debates += cmv
-        print(f"  CMV: {len(cmv)} debates  (total: {len(debates)})")
-    except Exception:
-        print("  CMV not found — training on IBM only")
+        print(f"  CMV (file): {len(cmv)} debates  (total: {len(debates)})")
+    except FileNotFoundError:
+        try:
+            from data import scrape_cmv
+            print("  CMV file missing — scraping live from Reddit (limit=150)…")
+            cmv = clean_debates(scrape_cmv(limit=150, sort="top", time_filter="all"))
+            if cmv:
+                debates += cmv
+                print(f"  CMV (Reddit live): {len(cmv)} debates  (total: {len(debates)})")
+            else:
+                print("  Reddit scraper returned 0 debates — training on IBM only")
+        except Exception as e2:
+            print(f"  CMV not available ({e2.__class__.__name__}: {e2}) — training on IBM only")
+    except Exception as e:
+        print(f"  CMV load error ({e}) — training on IBM only")
+
+    # Heavy imports after data loading so torch/CUDA init
+    # doesn't compete with dataset memory usage
+    import torch
+    from torch.utils.data import DataLoader, random_split
+    from torch.optim import AdamW
+    from transformers import (
+        RobertaTokenizerFast,
+        get_linear_schedule_with_warmup,
+    )
+    from tqdm import tqdm
+    from .dataset import ArgumentDataset
+    from .model import MODEL_NAME, build_model
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Device: {device}")
 
     tokenizer = RobertaTokenizerFast.from_pretrained(MODEL_NAME)
     dataset = ArgumentDataset(debates, tokenizer, max_length)
@@ -84,11 +109,15 @@ def train(
         [len(dataset) - val_size, val_size],
         generator=torch.Generator().manual_seed(42),
     )
-    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0)
-    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=0)
+    train_loader = DataLoader(
+        train_set, batch_size=batch_size, shuffle=True, num_workers=0,
+    )
+    val_loader = DataLoader(
+        val_set, batch_size=batch_size, shuffle=False, num_workers=0,
+    )
 
-    model = build_model().to(DEVICE)
-    weights = _class_weights(dataset).to(DEVICE)
+    model = build_model().to(device)
+    weights = _class_weights(dataset, device)
     loss_fn = torch.nn.CrossEntropyLoss(weight=weights)
 
     optimizer = AdamW(model.parameters(), lr=lr, weight_decay=0.01)
@@ -106,12 +135,15 @@ def train(
         model.train()
         total_loss = 0.0
         for batch in tqdm(train_loader, desc=f"Epoch {epoch}/{epochs}"):
-            input_ids = batch["input_ids"].to(DEVICE)
-            attention_mask = batch["attention_mask"].to(DEVICE)
-            labels = batch["label"].to(DEVICE)
+            input_ids = batch["input_ids"].to(device)
+            attention_mask = batch["attention_mask"].to(device)
+            labels = batch["label"].to(device)
 
             optimizer.zero_grad()
-            logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
+            logits = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            ).logits
             loss = loss_fn(logits, labels)
             loss.backward()
             torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
@@ -119,7 +151,7 @@ def train(
             scheduler.step()
             total_loss += loss.item()
 
-        val_loss, val_acc, val_f1 = evaluate(model, val_loader)
+        val_loss, val_acc, val_f1 = evaluate(model, val_loader, device)
         print(
             f"Epoch {epoch}:  "
             f"train_loss={total_loss / len(train_loader):.4f}  "
@@ -132,7 +164,10 @@ def train(
             best_f1 = val_f1
             model.save_pretrained(best_ckpt)
             tokenizer.save_pretrained(best_ckpt)
-            print(f"  → Saved best checkpoint (F1={best_f1:.4f}) to {best_ckpt}/")
+            print(
+                f"  → Saved best checkpoint "
+                f"(F1={best_f1:.4f}) to {best_ckpt}/"
+            )
 
     with open(os.path.join(output_dir, "train_config.json"), "w") as f:
         json.dump(
@@ -152,7 +187,7 @@ def train(
 
 
 if __name__ == "__main__":
-    p = argparse.ArgumentParser(description="Train argument classifier. Run as: python -m src.train")
+    p = argparse.ArgumentParser()
     p.add_argument("--epochs", type=int, default=3)
     p.add_argument("--batch-size", type=int, default=16)
     p.add_argument("--lr", type=float, default=2e-5)