Upload 3 files

SRL_preprocessing.py

@@ -0,0 +1,209 @@
+import argparse, json, torch
+from dataclasses import dataclass
+from typing import List, Dict, Optional
+from transformers import AutoTokenizer
+from torch.utils.data import Dataset
+
+# ==============================================================
+# 1. Data structure
+# ==============================================================
+@dataclass
+class SRLSample:
+    words: List[str]
+    predicate_word_idx: int
+    labels: List[str]
+    predicate_form: Optional[str] = None
+
+
+# ==============================================================
+# 2. Bracket → BIO conversion (unchanged)
+# ==============================================================
+def _bio_from_brackets(tags):
+    bio, stack = [], []
+    for t in tags:
+        if "(V*" in t:
+            bio.append("B-V")
+            continue
+        opens = []
+        i = 0
+        while True:
+            s = t.find("(", i)
+            if s == -1: break
+            e = t.find("*", s)
+            if e == -1: break
+            role = t[s+1:e]
+            opens.append(role)
+            i = e + 1
+        closes = t.count(")")
+        if opens:
+            bio.append(f"B-{opens[0]}")
+            for r in opens: stack.append(r)
+        elif stack:
+            bio.append(f"I-{stack[-1]}")
+        else:
+            bio.append("O")
+        for _ in range(closes):
+            if stack: stack.pop()
+    return bio
+
+
+def _read_sentences(path):
+    sent = []
+    with open(path, "r", encoding="utf-8", errors="replace") as f:
+        for line in f:
+            line = line.rstrip("\n")
+            if not line:
+                if sent: yield sent; sent = []
+                continue
+            cols = line.split()
+            if cols: sent.append(cols)
+    if sent: yield sent
+
+
+# ==============================================================
+# 3. CoNLL → SRLSample objects (in-memory)
+# ==============================================================
+def load_conll_samples(in_path, word_col_idx=3, srl_first_col_idx=11):
+    """
+    Reads .gold_conll file and returns list[SRLSample],
+    one per predicate column.
+    """
+    samples = []
+    for sent in _read_sentences(in_path):
+        words = [row[word_col_idx] for row in sent]
+        max_cols = max(len(row) for row in sent)
+        for srl_col in range(srl_first_col_idx, max_cols):
+            tags = [row[srl_col] if srl_col < len(row) else "*" for row in sent]
+            try:
+                pred_idx = next(i for i, t in enumerate(tags) if "(V*" in t)
+            except StopIteration:
+                continue
+            labels = _bio_from_brackets(tags)
+            predicate_form = words[pred_idx]
+            samples.append(SRLSample(words, pred_idx, labels, predicate_form))
+    print(f"[SRL_preprocessing] Loaded {len(samples)} predicate instances from {in_path}")
+    return samples
+
+
+# ==============================================================
+# 4. Dataset + Collate (same as yours, lightly cleaned)
+# ==============================================================
+class SRLDataset(Dataset):
+    def __init__(self, samples: List[SRLSample], tokenizer: AutoTokenizer,
+                 label2id: Dict[str, int], max_length: int = 256, debug_print=False):
+        self.samples = samples
+        self.tokenizer = tokenizer
+        self.label2id = label2id
+        self.id2label = {v: k for k, v in label2id.items()}
+        self.max_length = max_length
+        self.debug_print = debug_print
+
+    def __len__(self): return len(self.samples)
+
+    def _tokenize_sentence(self, words):
+        return self.tokenizer(words, is_split_into_words=True,
+                              add_special_tokens=False, return_attention_mask=False,
+                              return_token_type_ids=False)
+
+    def _tokenize_predicate(self, form):
+        return self.tokenizer(form, add_special_tokens=False,
+                              return_attention_mask=False,
+                              return_token_type_ids=False)
+
+    def __getitem__(self, idx):
+        instance = self.samples[idx]
+        words = instance.words
+        n_words = len(words)
+        pred_form = instance.predicate_form or words[instance.predicate_word_idx]
+
+        enc_sent = self._tokenize_sentence(words)
+        enc_pred = self._tokenize_predicate(pred_form)
+        sent_wp_ids = enc_sent["input_ids"]
+        pred_wp_ids = enc_pred["input_ids"]
+
+        input_ids = [self.tokenizer.cls_token_id] + sent_wp_ids + [self.tokenizer.sep_token_id] \
+                    + pred_wp_ids + [self.tokenizer.sep_token_id]
+        ttids = [0] * (1 + len(sent_wp_ids) + 1) + [1] * (len(pred_wp_ids) + 1)
+
+        tmp = self.tokenizer(words, is_split_into_words=True)
+        word_ids = tmp.word_ids()
+        first_pos_by_wid = {}
+        for pos, wid in enumerate(word_ids):
+            if wid is not None and wid not in first_pos_by_wid:
+                first_pos_by_wid[wid] = pos
+        word_first_wp_fullidx = [first_pos_by_wid[w] for w in range(n_words)]
+
+        label_ids = [self.label2id[l] for l in instance.labels]
+        indicator = [0]*n_words; indicator[instance.predicate_word_idx] = 1
+        attention_mask = [1]*len(input_ids)
+
+        if len(input_ids) > self.max_length:
+            max_pos = self.max_length-1
+            input_ids = input_ids[:self.max_length]
+            ttids = ttids[:self.max_length]
+            attention_mask = attention_mask[:self.max_length]
+            word_first_wp_fullidx = [min(p, max_pos) for p in word_first_wp_fullidx]
+
+        return {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+            "token_type_ids": torch.tensor(ttids, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+            "word_first_wp_fullidx": torch.tensor(word_first_wp_fullidx, dtype=torch.long),
+            "labels": torch.tensor(label_ids, dtype=torch.long),
+            "indicator": torch.tensor(indicator, dtype=torch.long),
+            "sent_len": torch.tensor(len(words), dtype=torch.long),
+            "pred_word_idx": torch.tensor(instance.predicate_word_idx, dtype=torch.long)
+        }
+
+
+def srl_collate(batch: List[Dict], pad_token_id: int, pad_label_id: int = -100):
+    B = len(batch)
+    max_L = max(item["input_ids"].size(0) for item in batch)
+    input_ids = torch.full((B, max_L), pad_token_id, dtype=torch.long)
+    token_type_ids = torch.zeros((B, max_L), dtype=torch.long)
+    attention_mask = torch.zeros((B, max_L), dtype=torch.long)
+    max_n = max(int(item["sent_len"]) for item in batch)
+    word_first_wp_fullidx = torch.full((B, max_n), -1, dtype=torch.long)
+    labels = torch.full((B, max_n), pad_label_id, dtype=torch.long)
+    indicator = torch.zeros((B, max_n), dtype=torch.long)
+    sent_lens = torch.zeros((B,), dtype=torch.long)
+    pred_word_idx = torch.zeros((B,), dtype=torch.long)
+    sentence_mask = torch.zeros((B, max_n), dtype=torch.bool)
+
+    for i, item in enumerate(batch):
+        L = item["input_ids"].size(0)
+        input_ids[i, :L] = item["input_ids"]
+        token_type_ids[i, :L] = item["token_type_ids"]
+        attention_mask[i, :L] = item["attention_mask"]
+        n = int(item["sent_len"])
+        word_first_wp_fullidx[i, :n] = item["word_first_wp_fullidx"]
+        labels[i, :n] = item["labels"]
+        indicator[i, :n] = item["indicator"]
+        sent_lens[i] = n
+        pred_word_idx[i] = item["pred_word_idx"]
+        sentence_mask[i, :n] = True
+
+    return {
+        "input_ids": input_ids,
+        "token_type_ids": token_type_ids,
+        "attention_mask": attention_mask,
+        "word_first_wp_fullidx": word_first_wp_fullidx,
+        "sentence_mask": sentence_mask,
+        "labels": labels,
+        "indicator": indicator,
+        "sent_lens": sent_lens,
+        "pred_word_idx": pred_word_idx,
+    }
+
+
+# ==============================================================
+# 5. Helper for trainer
+# ==============================================================
+def data_processing_for_loader_from_conll(conll_path, tokenizer, word_col_idx=3, srl_first_col_idx=11):
+    """Reads .gold_conll → SRLSample → SRLDataset, builds label2id/id2label."""
+    samples = load_conll_samples(conll_path, word_col_idx, srl_first_col_idx)
+    label_set = set(l for s in samples for l in s.labels)
+    label2id = {l: i for i, l in enumerate(sorted(label_set))}
+    id2label = {i: l for l, i in label2id.items()}
+    dataset = SRLDataset(samples, tokenizer, label2id, max_length=256)
+    return dataset, label2id, id2label

trainer.py

@@ -0,0 +1,288 @@
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, get_linear_schedule_with_warmup
+from sklearn.metrics import f1_score
+from torch.utils.data import DataLoader
+
+from SRL_preprocessing import data_processing_for_loader_from_conll, srl_collate
+from model import PredicateAwareSRL
+from utils import save_pkl
+import re, pathlib, argparse, json, os, sys
+
+
+try:
+    import _jsonnet
+except ImportError:
+    _jsonnet = None
+
+def load_cfg_from_jsonnet():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True, help="Path to .jsonnet config")
+    parser.add_argument("--out_dir", default=None, help="Override training.out_dir")
+    parser.add_argument("--best_model_path", default=None, help="Override best model save path")
+    parser.add_argument("--save_history_path", default=None, help="Override history pickle path")
+    args, unknown = parser.parse_known_args()
+
+    if _jsonnet is None:
+        raise RuntimeError("Please `pip install jsonnet` to use --config")
+
+    cfg = json.loads(_jsonnet.evaluate_file(args.config))
+
+    # Apply CLI overrides
+    if args.out_dir:
+        cfg.setdefault("training", {})["out_dir"] = args.out_dir
+
+    # Ensure out_dir exists & derive default file paths if missing
+    out_dir = cfg["training"].get("out_dir", "./checkpoints")
+    os.makedirs(out_dir, exist_ok=True)
+
+    # Derive defaults if not provided in config
+    cfg["training"].setdefault("best_model_path", os.path.join(out_dir, "best_srl_fr.ckpt"))
+    cfg["training"].setdefault("save_history_path", os.path.join(out_dir, "loss_history_fr.pkl"))
+
+    # Allow explicit overrides
+    if args.best_model_path:
+        cfg["training"]["best_model_path"] = args.best_model_path
+    if args.save_history_path:
+        cfg["training"]["save_history_path"] = args.save_history_path
+
+    return cfg
+
+# ==============================================================
+# 1. Training Loop
+# ==============================================================
+def train_one_epoch(
+    model,
+    dataloader,
+    optimizer,
+    device="cuda",
+    scheduler=None,
+    grad_accum_steps=1,
+    amp=True,
+    max_grad_norm=1.0,
+):
+    model.train()
+    total_loss, n_steps = 0.0, 0
+
+    use_amp = amp and torch.cuda.is_available()
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+
+    optimizer.zero_grad(set_to_none=True)
+
+    for step, batch in enumerate(dataloader, 1):
+        batch = {k: v.to(device) if torch.is_tensor(v) else v for k, v in batch.items()}
+
+        with torch.cuda.amp.autocast(enabled=use_amp, dtype=torch.float16):
+            _, loss = model(**batch)  # model must return (logits, loss)
+
+        total_loss += float(loss.detach().item())
+        n_steps += 1
+        loss = loss / grad_accum_steps
+
+        if use_amp:
+            scaler.scale(loss).backward()
+        else:
+            loss.backward()
+
+        if step % grad_accum_steps == 0:
+            if use_amp:
+                scaler.unscale_(optimizer)
+            nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
+
+            if use_amp:
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                optimizer.step()
+
+            optimizer.zero_grad(set_to_none=True)
+
+            if scheduler is not None:
+                scheduler.step()
+
+    return total_loss / max(1, n_steps)
+
+
+# ==============================================================
+# 2. Evaluation Loop
+# ==============================================================
+@torch.no_grad()
+def eval_loss_and_token_f1(model, dataloader, id2label=None, device="cuda", average="micro"):
+    model.eval()
+    total_loss, n_batches = 0.0, 0
+    all_preds, all_golds = [], []
+
+    for batch in dataloader:
+        gold = batch["labels"]
+        mask = (gold != -100)
+        batch = {k: v.to(device) if torch.is_tensor(v) else v for k, v in batch.items()}
+
+        logits, loss = model(**batch)
+        total_loss += float(loss.item()); n_batches += 1
+
+        preds = logits.argmax(-1).cpu()
+        all_preds.extend(preds[mask].tolist())
+        all_golds.extend(gold[mask].tolist())
+
+    f1 = f1_score(all_golds, all_preds, average=average)
+    return total_loss / max(1, n_batches), f1
+
+
+# ==============================================================
+# 3. Flexible Model Loader (English → French transfer)
+# ==============================================================
+def load_model(
+    bert_name: str,
+    label2id,
+    resume_path: str = None,
+    replace_encoder_with: str = None,
+    **kwargs
+):
+    """
+    Creates a PredicateAwareSRL model.
+    - If resume_path is given: loads SRL weights (English model)
+    - If replace_encoder_with is given: replaces only the BERT encoder
+      (e.g., replace 'bert-base-cased' with 'camembert-base')
+    """
+    print(f"🧩 Loading model backbone: {bert_name}")
+    model = PredicateAwareSRL(
+        bert_name=bert_name,
+        num_labels=len(label2id),
+        use_indicator=kwargs.get("use_indicator", True),
+        use_distance=kwargs.get("use_distance", True),
+        indicator_dim=kwargs.get("indicator_dim", 10),
+        lstm_hidden=kwargs.get("lstm_hidden", 768),
+        mlp_hidden=kwargs.get("mlp_hidden", 300),
+        pos_dim=kwargs.get("pos_dim", 50),
+        max_distance=kwargs.get("max_distance", 128),
+        dropout=kwargs.get("dropout", 0.1),
+        pad_label_id=-100,
+    )
+
+    if resume_path and os.path.exists(resume_path):
+        print(f"🔁 Loading SRL checkpoint from: {resume_path}")
+        state = torch.load(resume_path, map_location="cpu")
+        state_dict = state.get("model_state", state)
+        missing, unexpected = model.load_state_dict(state_dict, strict=False)
+        print(f"  → missing: {len(missing)}, unexpected: {len(unexpected)}")
+
+    if replace_encoder_with:
+        print(f"🌍 Replacing encoder with: {replace_encoder_with}")
+        from transformers import AutoModel
+        model.encoder = AutoModel.from_pretrained(replace_encoder_with)
+
+    return model
+
+
+# ==============================================================
+# 4. Main
+# ==============================================================
+if __name__ == "__main__":
+    # ------------------------------
+    # ⚙️ Configuration
+    # ------------------------------
+    cfg = load_cfg_from_jsonnet()
+
+    # read values from cfg as usual:
+    conll_train_path = cfg["data"]["conll_train"]
+    conll_dev_path   = cfg["data"].get("conll_dev")
+    word_col_idx     = cfg["data"]["word_col_idx"]
+    srl_first_col_idx= cfg["data"]["srl_first_col_idx"]
+
+    bert_name            = cfg["model"]["bert_name"]
+    resume_from          = cfg["model"].get("resume_from")
+    replace_encoder_with = cfg["model"].get("replace_encoder_with")
+    tok_name = (cfg["model"].get("tokenizer", {}) or {}).get("name", replace_encoder_with or bert_name)
+
+    out_dir         = cfg["training"]["out_dir"]
+    num_epochs      = cfg["training"]["num_epochs"]
+    batch_size      = cfg["training"]["batch_size"]
+    lr              = cfg["training"]["lr"]
+    weight_decay    = cfg["training"]["weight_decay"]
+    grad_accum      = cfg["training"]["grad_accum_steps"]
+    warmup_ratio    = cfg["training"]["warmup_ratio"]
+    amp             = cfg["training"]["amp"]
+    max_grad_norm   = cfg["training"]["max_grad_norm"]
+
+    best_model_path    = cfg["training"]["best_model_path"]
+    save_history_path  = cfg["training"]["save_history_path"]
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # ------------------------------
+    # 🧩 Tokenizer + data loading
+    # ------------------------------
+    tokenizer = AutoTokenizer.from_pretrained(replace_encoder_with or bert_name)
+    print(f"🔤 Using tokenizer: {replace_encoder_with or bert_name}")
+
+    print(f"📂 Loading French CoNLL data: {conll_train_path}")
+    train_dataset, label2id, id2label = data_processing_for_loader_from_conll(
+        conll_path=conll_train_path,
+        tokenizer=tokenizer,
+        word_col_idx=3,
+        srl_first_col_idx=11,
+    )
+
+    pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id
+    collate = lambda b: srl_collate(b, pad_token_id=pad_token_id, pad_label_id=-100)
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate)
+    dev_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate)  # no dev split yet
+
+    # ------------------------------
+    # 🧠 Model initialization
+    # ------------------------------
+    model = load_model(
+        bert_name=bert_name,
+        label2id=label2id,
+        resume_path=resume_from,
+        replace_encoder_with=replace_encoder_with,
+        use_indicator=True,
+        use_distance=True,
+        indicator_dim=10,
+        lstm_hidden=768,
+        mlp_hidden=300,
+        pos_dim=50,
+        max_distance=128,
+        dropout=0.1,
+    ).to(device)
+
+    # ------------------------------
+    # 🔧 Optimizer + Scheduler
+    # ------------------------------
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+    total_steps = len(train_loader) * num_epochs // max(1, grad_accum_steps)
+    warmup_steps = int(warmup_ratio * total_steps)
+
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=warmup_steps,
+        num_training_steps=total_steps,
+    )
+
+    # ------------------------------
+    # 🏋️ Training Loop
+    # ------------------------------
+    history = {"epoch": [], "train_loss": [], "dev_loss": [], "dev_f1": []}
+    best_dev, best_path = -1.0, "best_srl_fr.ckpt"
+
+    for epoch in range(num_epochs):
+        tr_loss = train_one_epoch(
+            model, train_loader, optimizer, device=device,
+            scheduler=scheduler, grad_accum_steps=grad_accum,
+            amp=True, max_grad_norm=1.0,
+        )
+        dev_loss, dev_f1 = eval_loss_and_token_f1(model, dev_loader, id2label, device=device)
+
+        history["epoch"].append(epoch + 1)
+        history["train_loss"].append(tr_loss)
+        history["dev_loss"].append(dev_loss)
+        history["dev_f1"].append(dev_f1)
+
+        print(f"Epoch {epoch+1}: train_loss={tr_loss:.4f}  dev_loss={dev_loss:.4f}  dev_F1={dev_f1:.4f}")
+
+        if dev_f1 > best_dev:
+            best_dev = dev_f1
+            torch.save({"model_state": model.state_dict(), "label2id": label2id}, best_path)
+            print(f"  ↳ new best dev; saved to {best_path}")
+
+    save_pkl(history, "loss_history_fr.pkl")

utils.py

@@ -0,0 +1,10 @@
+import pickle
+
+def save_pkl(tgt_list, svg_path):
+    with open(svg_path, "wb") as f:
+        pickle.dump(tgt_list, f)
+
+def load_pkl(path) :
+    with open(path, "rb") as f:
+        data = pickle.load(f)
+    return data