models and inference classes

src/models/dataset.py

@@ -0,0 +1,180 @@
+import json
+import numpy as np
+import torch
+from sklearn.model_selection import train_test_split
+from torch.utils.data import Dataset
+from src.schemas.labels import SENTIMENT_LABELS
+
+
+def load_data(path: str) -> list[dict]:
+    with open(path) as f:
+        return [json.loads(line) for line in f]
+
+
+def deduplicate_positions(samples: list[dict]) -> list[dict]:
+    """Select one position per entity.
+
+    Prefers the position whose position_text matches entity_text exactly
+    (case-insensitive). If none matches, selects the longest position.
+    """
+    out = []
+    for s in samples:
+        new_entities = []
+        for e in s["entities"]:
+            positions = e["positions"]
+            if not positions:
+                new_entities.append(e)
+                continue
+
+            exact = [
+                p for p in positions
+                if p["position_text"].lower() == e["entity_text"].lower()
+            ]
+
+            if exact:
+                best = max(exact, key=lambda p: p["length"])
+            else:
+                best = max(positions, key=lambda p: p["length"])
+
+            new_entities.append({**e, "positions": [best]})
+        out.append({**s, "entities": new_entities})
+    return out
+
+
+def flatten_to_examples(
+    samples: list[dict],
+    mode: str,
+) -> list[dict]:
+    """Flatten augmented data to one example per (entity, position) pair.
+
+    Reads pre-computed fields from the augmented JSONL:
+      marker -> seg_a = marker_text, seg_b = None
+      qa_m   -> seg_a = entity_centered_window, seg_b = qa_m_question
+      qa_b   -> 3 binary examples per position using qa_b_hypotheses
+    """
+    sentiments = list(SENTIMENT_LABELS.classes)
+    label2id = SENTIMENT_LABELS.label2id
+    examples = []
+
+    for s in samples:
+        for e in s["entities"]:
+            label_str = e.get("label")
+
+            base = {
+                "sample_id": s["id"],
+                "entity_id": e["entity_id"],
+                "entity_text": e["entity_text"],
+                "entity_type": e["entity_type"],
+            }
+
+            for p in e["positions"]:
+                if mode == "marker":
+                    ex = {**base, "seg_a": p["marker_text"], "seg_b": None}
+                    if label_str in label2id:
+                        ex["label"] = label2id[label_str]
+                    examples.append(ex)
+
+                elif mode == "qa_m":
+                    ex = {
+                        **base,
+                        "seg_a": p["entity_centered_window"],
+                        "seg_b": p["qa_m_question"],
+                    }
+                    if label_str in label2id:
+                        ex["label"] = label2id[label_str]
+                    examples.append(ex)
+
+                elif mode == "qa_b":
+                    for sentiment in sentiments:
+                        ex = {
+                            **base,
+                            "seg_a": p["entity_centered_window"],
+                            "seg_b": p["qa_b_hypotheses"][sentiment],
+                            "sentiment": sentiment,
+                        }
+                        if label_str in label2id:
+                            ex["label"] = 1 if sentiment == label_str else 0
+                        examples.append(ex)
+
+                else:
+                    raise ValueError(f"Unknown mode: {mode!r}")
+
+    return examples
+
+
+def split_data(
+    examples: list[dict], val_frac: float, test_frac: float, seed: int = 42
+) -> tuple[list[dict], list[dict], list[dict]]:
+    """Split at the *sample* level"""
+    sample_ids = np.array(list({e["sample_id"] for e in examples}))
+
+    remaining_ids, test_ids = train_test_split(
+        sample_ids, test_size=test_frac, random_state=seed
+    )
+    val_frac_adj = val_frac / (1.0 - test_frac)
+    train_ids, val_ids = train_test_split(
+        remaining_ids, test_size=val_frac_adj, random_state=seed
+    )
+
+    train_set = set(train_ids)
+    val_set = set(val_ids)
+    test_set = set(test_ids)
+
+    return (
+        [e for e in examples if e["sample_id"] in train_set],
+        [e for e in examples if e["sample_id"] in val_set],
+        [e for e in examples if e["sample_id"] in test_set],
+    )
+
+
+class EntitySentimentDataset(Dataset):
+    def __init__(self, examples: list[dict], tokenizer, max_len: int):
+        self.examples = examples
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, idx: int) -> dict:
+        ex = self.examples[idx]
+        seg_a = ex["seg_a"]
+        seg_b = ex["seg_b"]
+
+        if seg_b is None:
+            enc = self.tokenizer(
+                seg_a,
+                max_length=self.max_len,
+                truncation=True,
+                padding="max_length",
+                return_tensors="pt",
+            )
+        else:
+            enc = self.tokenizer(
+                seg_a, seg_b,
+                max_length=self.max_len,
+                truncation="only_first",
+                padding="max_length",
+                return_tensors="pt",
+            )
+
+        item = {
+            "input_ids": enc["input_ids"].squeeze(0),
+            "attention_mask": enc["attention_mask"].squeeze(0),
+        }
+        if "label" in ex:
+            item["labels"] = torch.tensor(ex["label"], dtype=torch.long)
+        return item
+
+
+class DeduplicatedEntitySentimentDataset(EntitySentimentDataset):
+    """Like EntitySentimentDataset but with one position per entity.
+
+    Applies deduplicate_positions before flattening, so each entity
+    contributes exactly one training example.
+    """
+
+    def __init__(self, samples: list[dict], mode: str, tokenizer, max_len: int):
+        deduped = deduplicate_positions(samples)
+        examples = flatten_to_examples(deduped, mode=mode)
+        super().__init__(examples, tokenizer, max_len)

src/models/distillbert.py

@@ -0,0 +1,127 @@
+from collections import Counter
+import numpy as np
+import torch
+import torch.nn.functional as F
+from sklearn.metrics import f1_score
+from torch import nn
+from torch.utils.data import DataLoader
+from transformers import Trainer
+from src.models.dataset import EntitySentimentDataset
+
+
+def compute_class_weights(examples: list[dict], n_classes: int) -> torch.Tensor:
+    counts = Counter(e["label"] for e in examples)
+    total = sum(counts.values())
+    weights = [total / (n_classes * counts.get(i, 1)) for i in range(n_classes)]
+    return torch.tensor(weights, dtype=torch.float)
+
+
+def focal_loss(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    weight: torch.Tensor,
+    gamma: float = 2.0,
+) -> torch.Tensor:
+    ce = F.cross_entropy(logits, labels, weight=weight, reduction="none")
+    probs = F.softmax(logits, dim=-1)
+    pt = probs.gather(1, labels.unsqueeze(1)).squeeze(1)
+    return ((1 - pt) ** gamma * ce).mean()
+
+
+class WeightedLossTrainer(Trainer):
+
+    def __init__(self, *args, class_weights: torch.Tensor, loss_fn: str = "cross_entropy", focal_gamma: float = 2.0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.class_weights = class_weights
+        self.loss_fn = loss_fn
+        self.focal_gamma = focal_gamma
+
+    def compute_loss(self, model, inputs, return_outputs: bool = False, **kwargs):
+        labels = inputs.pop("labels")
+        outputs = model(**inputs)
+        w = self.class_weights.to(outputs.logits.device)
+        if self.loss_fn == "focal":
+            loss = focal_loss(outputs.logits, labels, weight=w, gamma=self.focal_gamma)
+        else:
+            loss = nn.CrossEntropyLoss(weight=w)(outputs.logits, labels)
+        return (loss, outputs) if return_outputs else loss
+
+
+def reconstruct_triplets(
+    yes_probs: np.ndarray, bin_labels: np.ndarray
+) -> tuple[list[int], list[int]]:
+    """Group consecutive (neg, neu, pos) triplets and take argmax."""
+    preds3, labels3 = [], []
+    for i in range(0, len(yes_probs) - 2, 3):
+        preds3.append(int(np.argmax(yes_probs[i: i + 3])))
+        labels3.append(int(np.argmax(bin_labels[i: i + 3])))
+    return preds3, labels3
+
+
+def make_compute_metrics(mode: str):
+    if mode in ("marker", "qa_m"):
+        def compute_metrics(eval_pred):
+            logits, labels = eval_pred
+            preds = np.argmax(logits, axis=-1)
+            macro_f1 = f1_score(labels, preds, average="macro")
+            per_class = f1_score(labels, preds, average=None, labels=[0, 1, 2])
+            return {
+                "macro_f1": macro_f1,
+                "f1_negative": per_class[0],
+                "f1_neutral": per_class[1],
+                "f1_positive": per_class[2],
+            }
+    else:
+        def compute_metrics(eval_pred):
+            logits, labels = eval_pred
+            preds = np.argmax(logits, axis=-1)
+            bin_acc = float((preds == labels).mean())
+            bin_f1 = float(f1_score(labels, preds, average="binary", pos_label=1))
+
+            yes_probs = F.softmax(
+                torch.tensor(logits, dtype=torch.float), dim=-1
+            )[:, 1].numpy()
+
+            preds3, labels3 = reconstruct_triplets(yes_probs, labels)
+
+            macro_f1 = float(f1_score(preds3, labels3, average="macro")) \
+                if preds3 else 0.0
+            return {
+                "macro_f1": macro_f1,
+                "bin_accuracy": bin_acc,
+                "bin_f1_yes": bin_f1,
+            }
+
+    return compute_metrics
+
+
+def evaluate_qa_b_test(
+    model,
+    tokenizer,
+    test_exs: list[dict],
+    max_len: int,
+    batch_size: int,
+    device: torch.device,
+) -> tuple[float, list[int], list[int]]:
+    ds = EntitySentimentDataset(test_exs, tokenizer, max_len)
+    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)
+
+    all_yes_probs, all_bin_labels = [], []
+    model.eval()
+    with torch.no_grad():
+        for batch in loader:
+            logits = model(
+                input_ids=batch["input_ids"].to(device),
+                attention_mask=batch["attention_mask"].to(device),
+            ).logits
+            all_yes_probs.extend(
+                F.softmax(logits, dim=-1)[:, 1].cpu().tolist()
+            )
+            all_bin_labels.extend(batch["labels"].tolist())
+
+    preds3, labels3 = reconstruct_triplets(
+        np.array(all_yes_probs), np.array(all_bin_labels)
+    )
+
+    macro_f1 = f1_score(labels3, preds3, average="macro")
+    return macro_f1, preds3, labels3

src/models/fasttext.py

@@ -0,0 +1,197 @@
+import argparse
+import json
+import tempfile
+from pathlib import Path
+import fasttext
+import numpy as np
+from sklearn.metrics import f1_score, classification_report
+from src.models.augment import augment, MAXLEN_TO_WINDOW
+from src.models.dataset import deduplicate_positions, flatten_to_examples, split_data
+from src.schemas.labels import SENTIMENT_LABELS
+
+MODE = "marker"
+LABEL_PREFIX = "__label__"
+
+
+def _to_fasttext_line(example: dict) -> str:
+    text = example["seg_a"].replace("\n", " ")
+    label = SENTIMENT_LABELS.id2label[example["label"]]
+    return f"{LABEL_PREFIX}{label} {text}"
+
+
+def _write_fasttext_file(examples: list[dict], path: Path) -> None:
+    with open(path, "w", encoding="utf-8") as f:
+        for ex in examples:
+            f.write(_to_fasttext_line(ex) + "\n")
+
+
+def prepare_data(
+    data_path: str = "data/data_augmented_256.jsonl",
+    val_split: float = 0.1,
+    test_split: float = 0.1,
+    seed: int = 42,
+) -> tuple[list[dict], list[dict], list[dict]]:
+    with open(data_path, "r", encoding="utf-8") as f:
+        samples = [json.loads(line) for line in f]
+
+    examples = flatten_to_examples(samples, mode=MODE)
+    train_ex, val_ex, test_ex = split_data(examples, val_split, test_split, seed)
+
+    print(f"Train: {len(train_ex)}, Val: {len(val_ex)}, Test: {len(test_ex)}")
+    return train_ex, val_ex, test_ex
+
+
+def train(
+    train_examples: list[dict],
+    val_examples: list[dict],
+    output_dir: str = "models/fasttext",
+    lr: float = 0.5,
+    epoch: int = 25,
+    word_ngrams: int = 2,
+    dim: int = 100,
+    min_count: int = 1,
+) -> fasttext.FastText._FastText:
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    train_file = output_dir / "train.txt"
+    _write_fasttext_file(train_examples, train_file)
+
+    model = fasttext.train_supervised(
+        input=str(train_file),
+        lr=lr,
+        epoch=epoch,
+        wordNgrams=word_ngrams,
+        dim=dim,
+        minCount=min_count,
+        loss="softmax",
+    )
+
+    model.save_model(str(output_dir / "model.bin"))
+    print(f"Model saved to {output_dir / 'model.bin'}")
+
+    evaluate(model, val_examples, split_name="val")
+
+    return model
+
+
+def evaluate(
+    model: fasttext.FastText._FastText,
+    examples: list[dict],
+    split_name: str = "test",
+) -> float:
+    sentiments = list(SENTIMENT_LABELS.classes)
+    true_labels = []
+    pred_labels = []
+
+    for ex in examples:
+        text = ex["seg_a"].replace("\n", " ")
+        prediction = model.predict(text)[0][0].replace(LABEL_PREFIX, "")
+        pred_labels.append(prediction)
+        true_labels.append(SENTIMENT_LABELS.id2label[ex["label"]])
+
+    macro_f1 = f1_score(true_labels, pred_labels, average="macro", labels=sentiments)
+    print(f"\n{split_name} (per-position) macro F1: {macro_f1:.4f}")
+    print(classification_report(true_labels, pred_labels, labels=sentiments, digits=4))
+
+    return macro_f1
+
+
+def evaluate_entity_level(
+    model: fasttext.FastText._FastText,
+    examples: list[dict],
+    split_name: str = "test",
+) -> float:
+    sentiments = list(SENTIMENT_LABELS.classes)
+
+    entity_preds: dict[tuple, tuple[str, float]] = {}
+    entity_labels: dict[tuple, str] = {}
+
+    for ex in examples:
+        key = (ex["sample_id"], ex["entity_id"])
+        text = ex["seg_a"].replace("\n", " ")
+        labels, probs = model.predict(text)
+        label = labels[0].replace(LABEL_PREFIX, "")
+        conf = float(probs[0])
+        if key not in entity_preds or conf > entity_preds[key][1]:
+            entity_preds[key] = (label, conf)
+        entity_labels[key] = SENTIMENT_LABELS.id2label[ex["label"]]
+
+    true = [entity_labels[k] for k in entity_preds]
+    pred = [entity_preds[k][0] for k in entity_preds]
+
+    macro_f1 = f1_score(true, pred, average="macro", labels=sentiments)
+    print(f"\n{split_name} (entity-level) macro F1: {macro_f1:.4f}")
+    print(classification_report(true, pred, labels=sentiments, digits=4))
+
+    return macro_f1
+
+
+def predict_samples(
+    model: fasttext.FastText._FastText,
+    samples: list[dict],
+    window_words: int = 70,
+    deduplicate: bool = False,
+) -> list[dict]:
+    augmented = augment(samples, window_words)
+    if deduplicate:
+        augmented = deduplicate_positions(augmented)
+    examples = flatten_to_examples(augmented, mode=MODE)
+
+    entity_preds: dict[tuple, tuple[str, float]] = {}
+    for ex in examples:
+        key = (ex["sample_id"], ex["entity_id"])
+        text = ex["seg_a"].replace("\n", " ")
+        labels, probs = model.predict(text)
+        label = labels[0].replace(LABEL_PREFIX, "")
+        conf = float(probs[0])
+        if key not in entity_preds or conf > entity_preds[key][1]:
+            entity_preds[key] = (label, conf)
+
+    results = []
+    for s in samples:
+        entities_out = []
+        for e in s["entities"]:
+            key = (s["id"], e["entity_id"])
+            entities_out.append({
+                "entity_id": e["entity_id"],
+                "entity_text": e["entity_text"],
+                "classification": entity_preds.get(key, ("neutral", 0.0))[0],
+            })
+        results.append({"id": s["id"], "entities": entities_out})
+
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="fastText baseline for entity sentiment")
+    parser.add_argument("--data", default="data/data_augmented_256.jsonl")
+    parser.add_argument("--output-dir", default="models/fasttext")
+    parser.add_argument("--lr", type=float, default=0.5)
+    parser.add_argument("--epoch", type=int, default=25)
+    parser.add_argument("--word-ngrams", type=int, default=2)
+    parser.add_argument("--dim", type=int, default=100)
+    parser.add_argument("--val-split", type=float, default=0.1)
+    parser.add_argument("--test-split", type=float, default=0.1)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    train_ex, val_ex, test_ex = prepare_data(
+        args.data, args.val_split, args.test_split, args.seed,
+    )
+
+    model = train(
+        train_ex, val_ex,
+        output_dir=args.output_dir,
+        lr=args.lr,
+        epoch=args.epoch,
+        word_ngrams=args.word_ngrams,
+        dim=args.dim,
+    )
+
+    evaluate(model, test_ex, split_name="test")
+    evaluate_entity_level(model, test_ex, split_name="test")
+
+
+if __name__ == "__main__":
+    main()

src/models/inference.py

@@ -0,0 +1,175 @@
+import argparse
+import json
+from pathlib import Path
+import numpy as np
+import onnxruntime as ort
+from transformers import AutoTokenizer
+from src.models.augment import augment, MAXLEN_TO_WINDOW
+from src.models.dataset import deduplicate_positions, flatten_to_examples
+from src.models.distillbert import reconstruct_triplets
+from src.schemas.labels import MARKER_MODE, SENTIMENT_LABELS
+
+
+BASE_TOKENIZER = "distilbert-base-uncased"
+
+def build_tokenizer(mode: str):
+    tokenizer = AutoTokenizer.from_pretrained(BASE_TOKENIZER)
+    if mode == "marker":
+        tokenizer.add_special_tokens(
+            {"additional_special_tokens": [MARKER_MODE.entity_start, MARKER_MODE.entity_end]}
+        )
+    return tokenizer
+
+
+def _softmax(logits: np.ndarray) -> np.ndarray:
+    exp = np.exp(logits - logits.max(axis=-1, keepdims=True))
+    return exp / exp.sum(axis=-1, keepdims=True)
+
+
+def _tokenize_examples(
+    examples: list[dict], tokenizer, max_len: int,
+) -> dict[str, np.ndarray]:
+    input_ids, attention_masks = [], []
+    for ex in examples:
+        seg_a = ex["seg_a"]
+        seg_b = ex["seg_b"]
+
+        if seg_b is None:
+            enc = tokenizer(
+                seg_a,
+                max_length=max_len,
+                truncation=True,
+                padding="max_length",
+                return_tensors="np",
+            )
+        else:
+            enc = tokenizer(
+                seg_a, seg_b,
+                max_length=max_len,
+                truncation="only_first",
+                padding="max_length",
+                return_tensors="np",
+            )
+        input_ids.append(enc["input_ids"][0])
+        attention_masks.append(enc["attention_mask"][0])
+
+    return {
+        "input_ids": np.array(input_ids, dtype=np.int64),
+        "attention_mask": np.array(attention_masks, dtype=np.int64),
+    }
+
+
+def _run_batched(
+    session: ort.InferenceSession,
+    inputs: dict[str, np.ndarray],
+    batch_size: int,
+) -> np.ndarray:
+    n = inputs["input_ids"].shape[0]
+    all_logits = []
+    for start in range(0, n, batch_size):
+        end = min(start + batch_size, n)
+        batch = {k: v[start:end] for k, v in inputs.items()}
+        logits = session.run(None, batch)[0]
+        all_logits.append(logits)
+    return np.concatenate(all_logits, axis=0)
+
+
+def predict(
+    samples: list[dict],
+    session: ort.InferenceSession,
+    tokenizer,
+    mode: str,
+    max_len: int = 256,
+    batch_size: int = 32,
+    deduplicate: bool = False,
+) -> list[dict]:
+    window_words = MAXLEN_TO_WINDOW[max_len]
+    augmented = augment(samples, window_words)
+    if deduplicate:
+        augmented = deduplicate_positions(augmented)
+    examples = flatten_to_examples(augmented, mode=mode)
+
+    if not examples:
+        return [{"id": s["id"], "entities": []} for s in samples]
+
+    inputs = _tokenize_examples(examples, tokenizer, max_len)
+    logits = _run_batched(session, inputs, batch_size)
+    sentiments = list(SENTIMENT_LABELS.classes)
+
+    probs = _softmax(logits)
+
+    if mode in ("marker", "qa_m"):
+        preds = np.argmax(probs, axis=-1)
+        max_probs = probs.max(axis=-1)
+        for ex, pred_id, conf in zip(examples, preds, max_probs):
+            ex["predicted_label"] = sentiments[int(pred_id)]
+            ex["confidence"] = float(conf)
+
+    else:
+        yes_probs = probs[:, 1]
+        preds3, _ = reconstruct_triplets(yes_probs, np.zeros_like(yes_probs))
+
+        triplet_idx = 0
+        i = 0
+        while i < len(examples) - 2:
+            pred_label = sentiments[preds3[triplet_idx]]
+            triplet_conf = float(yes_probs[i:i + 3].max())
+            for j in range(3):
+                examples[i + j]["predicted_label"] = pred_label
+                examples[i + j]["confidence"] = triplet_conf
+            triplet_idx += 1
+            i += 3
+
+    entity_preds: dict[tuple, tuple[str, float]] = {}
+    for ex in examples:
+        key = (ex["sample_id"], ex["entity_id"])
+        conf = ex.get("confidence", 0.0)
+        if key not in entity_preds or conf > entity_preds[key][1]:
+            entity_preds[key] = (ex["predicted_label"], conf)
+
+    results = []
+    for s in samples:
+        entities_out = []
+        for e in s["entities"]:
+            key = (s["id"], e["entity_id"])
+            label, _ = entity_preds.get(key, ("neutral", 0.0))
+            entities_out.append({
+                "entity_id": e["entity_id"],
+                "entity_text": e["entity_text"],
+                "classification": label,
+            })
+        results.append({"id": s["id"], "entities": entities_out})
+
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run ONNX inference on raw input JSON")
+    parser.add_argument("--onnx-path", required=True, help="Path to model.onnx")
+    parser.add_argument("--mode", required=True, choices=("marker", "qa_m", "qa_b"))
+    parser.add_argument("--data", required=True, help="Path to input JSON (assignment format)")
+    parser.add_argument("--output", default=None, help="Output JSON path (default: stdout)")
+    parser.add_argument("--max-len", type=int, default=256)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--deduplicate", action="store_true", help="Use one position per entity")
+    args = parser.parse_args()
+
+    mode = args.mode
+    tokenizer = build_tokenizer(mode)
+    session = ort.InferenceSession(args.onnx_path)
+
+    with open(args.data, "r", encoding="utf-8") as f:
+        samples = json.load(f)
+
+    results = predict(samples, session, tokenizer, mode, args.max_len, args.batch_size, deduplicate=args.deduplicate)
+
+    output_json = json.dumps(results, ensure_ascii=False, indent=2)
+    if args.output:
+        Path(args.output).write_text(output_json, encoding="utf-8")
+        print(f"Saved {len(results)} predictions to {args.output}")
+    else:
+        print(output_json)
+
+
+if __name__ == "__main__":
+    main()