Upload 7 files

backup/modules/base.py

@@ -0,0 +1,150 @@
+from collections import defaultdict
+from typing import List, Tuple, Dict
+
+import torch
+from torch import nn
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader
+import random
+
+
+class InstructBase(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.max_width = config.max_width
+        self.base_config = config
+
+    def get_dict(self, spans, classes_to_id):
+        dict_tag = defaultdict(int)
+        for span in spans:
+            if span[2] in classes_to_id:
+                dict_tag[(span[0], span[1])] = classes_to_id[span[2]]
+        return dict_tag
+
+    def preprocess_spans(self, tokens, ner, classes_to_id):
+
+        max_len = self.base_config.max_len
+
+        if len(tokens) > max_len:
+            length = max_len
+            tokens = tokens[:max_len]
+        else:
+            length = len(tokens)
+
+        spans_idx = []
+        for i in range(length):
+            spans_idx.extend([(i, i + j) for j in range(self.max_width)])
+
+        dict_lab = self.get_dict(ner, classes_to_id) if ner else defaultdict(int)
+
+        # 0 for null labels
+        span_label = torch.LongTensor([dict_lab[i] for i in spans_idx])
+        spans_idx = torch.LongTensor(spans_idx)
+
+        # mask for valid spans
+        valid_span_mask = spans_idx[:, 1] > length - 1
+
+        # mask invalid positions
+        span_label = span_label.masked_fill(valid_span_mask, -1)
+
+        return {
+            'tokens': tokens,
+            'span_idx': spans_idx,
+            'span_label': span_label,
+            'seq_length': length,
+            'entities': ner,
+        }
+
+    def collate_fn(self, batch_list, entity_types=None):
+        # batch_list: list of dict containing tokens, ner
+        if entity_types is None:
+            negs = self.get_negatives(batch_list, 100)
+            class_to_ids = []
+            id_to_classes = []
+            for b in batch_list:
+                # negs = b["negative"]
+                random.shuffle(negs)
+
+                # negs = negs[:sampled_neg]
+                max_neg_type_ratio = int(self.base_config.max_neg_type_ratio)
+
+                if max_neg_type_ratio == 0:
+                    # no negatives
+                    neg_type_ratio = 0
+                else:
+                    neg_type_ratio = random.randint(0, max_neg_type_ratio)
+
+                if neg_type_ratio == 0:
+                    # no negatives
+                    negs_i = []
+                else:
+                    negs_i = negs[:len(b['ner']) * neg_type_ratio]
+
+                # this is the list of all possible entity types (positive and negative)
+                types = list(set([el[-1] for el in b['ner']] + negs_i))
+
+                # shuffle (every epoch)
+                random.shuffle(types)
+
+                if len(types) != 0:
+                    # prob of higher number shoul
+                    # random drop
+                    if self.base_config.random_drop:
+                        num_ents = random.randint(1, len(types))
+                        types = types[:num_ents]
+
+                # maximum number of entities types
+                types = types[:int(self.base_config.max_types)]
+
+                # supervised training
+                if "label" in b:
+                    types = sorted(b["label"])
+
+                class_to_id = {k: v for v, k in enumerate(types, start=1)}
+                id_to_class = {k: v for v, k in class_to_id.items()}
+                class_to_ids.append(class_to_id)
+                id_to_classes.append(id_to_class)
+
+            batch = [
+                self.preprocess_spans(b["tokenized_text"], b["ner"], class_to_ids[i]) for i, b in enumerate(batch_list)
+            ]
+
+        else:
+            class_to_ids = {k: v for v, k in enumerate(entity_types, start=1)}
+            id_to_classes = {k: v for v, k in class_to_ids.items()}
+            batch = [
+                self.preprocess_spans(b["tokenized_text"], b["ner"], class_to_ids) for b in batch_list
+            ]
+
+        span_idx = pad_sequence(
+            [b['span_idx'] for b in batch], batch_first=True, padding_value=0
+        )
+
+        span_label = pad_sequence(
+            [el['span_label'] for el in batch], batch_first=True, padding_value=-1
+        )
+
+        return {
+            'seq_length': torch.LongTensor([el['seq_length'] for el in batch]),
+            'span_idx': span_idx,
+            'tokens': [el['tokens'] for el in batch],
+            'span_mask': span_label != -1,
+            'span_label': span_label,
+            'entities': [el['entities'] for el in batch],
+            'classes_to_id': class_to_ids,
+            'id_to_classes': id_to_classes,
+        }
+
+    @staticmethod
+    def get_negatives(batch_list, sampled_neg=5):
+        ent_types = []
+        for b in batch_list:
+            types = set([el[-1] for el in b['ner']])
+            ent_types.extend(list(types))
+        ent_types = list(set(ent_types))
+        # sample negatives
+        random.shuffle(ent_types)
+        return ent_types[:sampled_neg]
+
+    def create_dataloader(self, data, entity_types=None, **kwargs):
+        return DataLoader(data, collate_fn=lambda x: self.collate_fn(x, entity_types), **kwargs)

backup/modules/data_proc.py

@@ -0,0 +1,73 @@
+import json
+from tqdm import tqdm
+# ast.literal_eval
+import ast, re
+
+path = 'train.json'
+
+with open(path, 'r') as f:
+    data = json.load(f)
+
+def tokenize_text(text):
+    return re.findall(r'\w+(?:[-_]\w+)*|\S', text)
+
+def extract_entity_spans(entry):
+    text = ""
+    len_start = len("What describes ")
+    len_end = len(" in the text?")
+    entity_types = []
+    entity_texts = []
+
+    for c in entry['conversations']:
+        if c['from'] == 'human' and c['value'].startswith('Text: '):
+            text = c['value'][len('Text: '):]
+            tokenized_text = tokenize_text(text)
+
+        if c['from'] == 'human' and c['value'].startswith('What describes '):
+
+            c_type = c['value'][len_start:-len_end]
+            c_type = c_type.replace(' ', '_')
+            entity_types.append(c_type)
+
+        elif c['from'] == 'gpt' and c['value'].startswith('['):
+            if c['value'] == '[]':
+                entity_types = entity_types[:-1]
+                continue
+
+            texts_ents = ast.literal_eval(c['value'])
+            # replace space to _ in texts_ents
+            entity_texts.extend(texts_ents)
+            num_repeat = len(texts_ents) - 1
+            entity_types.extend([entity_types[-1]] * num_repeat)
+
+    entity_spans = []
+    for j, entity_text in enumerate(entity_texts):
+        entity_tokens = tokenize_text(entity_text)
+        matches = []
+        for i in range(len(tokenized_text) - len(entity_tokens) + 1):
+            if " ".join(tokenized_text[i:i + len(entity_tokens)]).lower() == " ".join(entity_tokens).lower():
+                matches.append((i, i + len(entity_tokens) - 1, entity_types[j]))
+        if matches:
+            entity_spans.extend(matches)
+
+    return entity_spans, tokenized_text
+
+# Usage:
+# Replace 'entry' with the specific entry from your JSON data
+entry = data[17818]  # For example, taking the first entry
+entity_spans, tokenized_text = extract_entity_spans(entry)
+print("Entity Spans:", entity_spans)
+#print("Tokenized Text:", tokenized_text)
+
+# create a dict: {"tokenized_text": tokenized_text, "entity_spans": entity_spans}
+
+all_data = []
+
+for entry in tqdm(data):
+    entity_spans, tokenized_text = extract_entity_spans(entry)
+    all_data.append({"tokenized_text": tokenized_text, "ner": entity_spans})
+
+
+with open('train_instruct.json', 'w') as f:
+    json.dump(all_data, f)
+

backup/modules/evaluator.py

@@ -0,0 +1,152 @@
+from collections import defaultdict
+
+import numpy as np
+import torch
+from seqeval.metrics.v1 import _prf_divide
+
+
+def extract_tp_actual_correct(y_true, y_pred):
+    entities_true = defaultdict(set)
+    entities_pred = defaultdict(set)
+
+    for type_name, (start, end), idx in y_true:
+        entities_true[type_name].add((start, end, idx))
+    for type_name, (start, end), idx in y_pred:
+        entities_pred[type_name].add((start, end, idx))
+
+    target_names = sorted(set(entities_true.keys()) | set(entities_pred.keys()))
+
+    tp_sum = np.array([], dtype=np.int32)
+    pred_sum = np.array([], dtype=np.int32)
+    true_sum = np.array([], dtype=np.int32)
+    for type_name in target_names:
+        entities_true_type = entities_true.get(type_name, set())
+        entities_pred_type = entities_pred.get(type_name, set())
+        tp_sum = np.append(tp_sum, len(entities_true_type & entities_pred_type))
+        pred_sum = np.append(pred_sum, len(entities_pred_type))
+        true_sum = np.append(true_sum, len(entities_true_type))
+
+    return pred_sum, tp_sum, true_sum, target_names
+
+
+def flatten_for_eval(y_true, y_pred):
+    all_true = []
+    all_pred = []
+
+    for i, (true, pred) in enumerate(zip(y_true, y_pred)):
+        all_true.extend([t + [i] for t in true])
+        all_pred.extend([p + [i] for p in pred])
+
+    return all_true, all_pred
+
+
+def compute_prf(y_true, y_pred, average='micro'):
+    y_true, y_pred = flatten_for_eval(y_true, y_pred)
+
+    pred_sum, tp_sum, true_sum, target_names = extract_tp_actual_correct(y_true, y_pred)
+
+    if average == 'micro':
+        tp_sum = np.array([tp_sum.sum()])
+        pred_sum = np.array([pred_sum.sum()])
+        true_sum = np.array([true_sum.sum()])
+
+    precision = _prf_divide(
+        numerator=tp_sum,
+        denominator=pred_sum,
+        metric='precision',
+        modifier='predicted',
+        average=average,
+        warn_for=('precision', 'recall', 'f-score'),
+        zero_division='warn'
+    )
+
+    recall = _prf_divide(
+        numerator=tp_sum,
+        denominator=true_sum,
+        metric='recall',
+        modifier='true',
+        average=average,
+        warn_for=('precision', 'recall', 'f-score'),
+        zero_division='warn'
+    )
+
+    denominator = precision + recall
+    denominator[denominator == 0.] = 1
+    f_score = 2 * (precision * recall) / denominator
+
+    return {'precision': precision[0], 'recall': recall[0], 'f_score': f_score[0]}
+
+
+class Evaluator:
+    def __init__(self, all_true, all_outs):
+        self.all_true = all_true
+        self.all_outs = all_outs
+
+    def get_entities_fr(self, ents):
+        all_ents = []
+        for s, e, lab in ents:
+            all_ents.append([lab, (s, e)])
+        return all_ents
+
+    def transform_data(self):
+        all_true_ent = []
+        all_outs_ent = []
+        for i, j in zip(self.all_true, self.all_outs):
+            e = self.get_entities_fr(i)
+            all_true_ent.append(e)
+            e = self.get_entities_fr(j)
+            all_outs_ent.append(e)
+        return all_true_ent, all_outs_ent
+
+    @torch.no_grad()
+    def evaluate(self):
+        all_true_typed, all_outs_typed = self.transform_data()
+        precision, recall, f1 = compute_prf(all_true_typed, all_outs_typed).values()
+        output_str = f"P: {precision:.2%}\tR: {recall:.2%}\tF1: {f1:.2%}\n"
+        return output_str, f1
+
+
+def is_nested(idx1, idx2):
+    # Return True if idx2 is nested inside idx1 or vice versa
+    return (idx1[0] <= idx2[0] and idx1[1] >= idx2[1]) or (idx2[0] <= idx1[0] and idx2[1] >= idx1[1])
+
+
+def has_overlapping(idx1, idx2):
+    overlapping = True
+    if idx1[:2] == idx2[:2]:
+        return overlapping
+    if (idx1[0] > idx2[1] or idx2[0] > idx1[1]):
+        overlapping = False
+    return overlapping
+
+
+def has_overlapping_nested(idx1, idx2):
+    # Return True if idx1 and idx2 overlap, but neither is nested inside the other
+    if idx1[:2] == idx2[:2]:
+        return True
+    if ((idx1[0] > idx2[1] or idx2[0] > idx1[1]) or is_nested(idx1, idx2)) and idx1 != idx2:
+        return False
+    else:
+        return True
+
+
+def greedy_search(spans, flat_ner=True):  # start, end, class, score
+
+    if flat_ner:
+        has_ov = has_overlapping
+    else:
+        has_ov = has_overlapping_nested
+
+    new_list = []
+    span_prob = sorted(spans, key=lambda x: -x[-1])
+    for i in range(len(spans)):
+        b = span_prob[i]
+        flag = False
+        for new in new_list:
+            if has_ov(b[:-1], new):
+                flag = True
+                break
+        if not flag:
+            new_list.append(b[:-1])
+    new_list = sorted(new_list, key=lambda x: x[0])
+    return new_list

backup/modules/layers.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
+
+
+class LstmSeq2SeqEncoder(nn.Module):
+    def __init__(self, input_size, hidden_size, num_layers=1, dropout=0., bidirectional=False):
+        super(LstmSeq2SeqEncoder, self).__init__()
+        self.lstm = nn.LSTM(input_size=input_size,
+                            hidden_size=hidden_size,
+                            num_layers=num_layers,
+                            dropout=dropout,
+                            bidirectional=bidirectional,
+                            batch_first=True)
+
+    def forward(self, x, mask, hidden=None):
+        # Packing the input sequence
+        lengths = mask.sum(dim=1).cpu()
+        packed_x = pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False)
+
+        # Passing packed sequence through LSTM
+        packed_output, hidden = self.lstm(packed_x, hidden)
+
+        # Unpacking the output sequence
+        output, _ = pad_packed_sequence(packed_output, batch_first=True)
+
+        return output

backup/modules/run_evaluation.py

@@ -0,0 +1,188 @@
+import glob
+import json
+import os
+import os
+
+import torch
+from tqdm import tqdm
+import random
+
+
+def open_content(path):
+    paths = glob.glob(os.path.join(path, "*.json"))
+    train, dev, test, labels = None, None, None, None
+    for p in paths:
+        if "train" in p:
+            with open(p, "r") as f:
+                train = json.load(f)
+        elif "dev" in p:
+            with open(p, "r") as f:
+                dev = json.load(f)
+        elif "test" in p:
+            with open(p, "r") as f:
+                test = json.load(f)
+        elif "labels" in p:
+            with open(p, "r") as f:
+                labels = json.load(f)
+    return train, dev, test, labels
+
+
+def process(data):
+    words = data['sentence'].split()
+    entities = []  # List of entities (start, end, type)
+
+    for entity in data['entities']:
+        start_char, end_char = entity['pos']
+
+        # Initialize variables to keep track of word positions
+        start_word = None
+        end_word = None
+
+        # Iterate through words and find the word positions
+        char_count = 0
+        for i, word in enumerate(words):
+            word_length = len(word)
+            if char_count == start_char:
+                start_word = i
+            if char_count + word_length == end_char:
+                end_word = i
+                break
+            char_count += word_length + 1  # Add 1 for the space
+
+        # Append the word positions to the list
+        entities.append((start_word, end_word, entity['type']))
+
+    # Create a list of word positions for each entity
+    sample = {
+        "tokenized_text": words,
+        "ner": entities
+    }
+
+    return sample
+
+
+# create dataset
+def create_dataset(path):
+    train, dev, test, labels = open_content(path)
+    train_dataset = []
+    dev_dataset = []
+    test_dataset = []
+    for data in train:
+        train_dataset.append(process(data))
+    for data in dev:
+        dev_dataset.append(process(data))
+    for data in test:
+        test_dataset.append(process(data))
+    return train_dataset, dev_dataset, test_dataset, labels
+
+
+@torch.no_grad()
+def get_for_one_path(path, model):
+    # load the dataset
+    _, _, test_dataset, entity_types = create_dataset(path)
+
+    data_name = path.split("/")[-1]  # get the name of the dataset
+
+    # check if the dataset is flat_ner
+    flat_ner = True
+    if any([i in data_name for i in ["ACE", "GENIA", "Corpus"]]):
+        flat_ner = False
+
+    # evaluate the model
+    results, f1 = model.evaluate(test_dataset, flat_ner=flat_ner, threshold=0.5, batch_size=12,
+                                 entity_types=entity_types)
+    return data_name, results, f1
+
+
+def get_for_all_path(model, steps, log_dir, data_paths):
+    all_paths = glob.glob(f"{data_paths}/*")
+
+    all_paths = sorted(all_paths)
+
+    # move the model to the device
+    device = next(model.parameters()).device
+    model.to(device)
+    # set the model to eval mode
+    model.eval()
+
+    # log the results
+    save_path = os.path.join(log_dir, "results.txt")
+
+    with open(save_path, "a") as f:
+        f.write("##############################################\n")
+        # write step
+        f.write("step: " + str(steps) + "\n")
+
+    zero_shot_benc = ["mit-movie", "mit-restaurant", "CrossNER_AI", "CrossNER_literature", "CrossNER_music",
+                      "CrossNER_politics", "CrossNER_science"]
+
+    zero_shot_benc_results = {}
+    all_results = {}  # without crossNER
+
+    for p in tqdm(all_paths):
+        if "sample_" not in p:
+            data_name, results, f1 = get_for_one_path(p, model)
+            # write to file
+            with open(save_path, "a") as f:
+                f.write(data_name + "\n")
+                f.write(str(results) + "\n")
+
+            if data_name in zero_shot_benc:
+                zero_shot_benc_results[data_name] = f1
+            else:
+                all_results[data_name] = f1
+
+    avg_all = sum(all_results.values()) / len(all_results)
+    avg_zs = sum(zero_shot_benc_results.values()) / len(zero_shot_benc_results)
+
+    save_path_table = os.path.join(log_dir, "tables.txt")
+
+    # results for all datasets except crossNER
+    table_bench_all = ""
+    for k, v in all_results.items():
+        table_bench_all += f"{k:20}: {v:.1%}\n"
+    # (20 size aswell for average i.e. :20)
+    table_bench_all += f"{'Average':20}: {avg_all:.1%}"
+
+    # results for zero-shot benchmark
+    table_bench_zeroshot = ""
+    for k, v in zero_shot_benc_results.items():
+        table_bench_zeroshot += f"{k:20}: {v:.1%}\n"
+    table_bench_zeroshot += f"{'Average':20}: {avg_zs:.1%}"
+
+    # write to file
+    with open(save_path_table, "a") as f:
+        f.write("##############################################\n")
+        f.write("step: " + str(steps) + "\n")
+        f.write("Table for all datasets except crossNER\n")
+        f.write(table_bench_all + "\n\n")
+        f.write("Table for zero-shot benchmark\n")
+        f.write(table_bench_zeroshot + "\n")
+        f.write("##############################################\n\n")
+
+
+def sample_train_data(data_paths, sample_size=10000):
+    all_paths = glob.glob(f"{data_paths}/*")
+
+    all_paths = sorted(all_paths)
+
+    # to exclude the zero-shot benchmark datasets
+    zero_shot_benc = ["CrossNER_AI", "CrossNER_literature", "CrossNER_music",
+                      "CrossNER_politics", "CrossNER_science", "ACE 2004"]
+
+    new_train = []
+    # take 10k samples from each dataset
+    for p in tqdm(all_paths):
+        if any([i in p for i in zero_shot_benc]):
+            continue
+        train, dev, test, labels = create_dataset(p)
+
+        # add label key to the train data
+        for i in range(len(train)):
+            train[i]["label"] = labels
+
+        random.shuffle(train)
+        train = train[:sample_size]
+        new_train.extend(train)
+
+    return new_train

backup/modules/span_rep.py

@@ -0,0 +1,369 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+def create_projection_layer(hidden_size: int, dropout: float, out_dim: int = None) -> nn.Sequential:
+    """
+    Creates a projection layer with specified configurations.
+    """
+    if out_dim is None:
+        out_dim = hidden_size
+
+    return nn.Sequential(
+        nn.Linear(hidden_size, out_dim * 4),
+        nn.ReLU(),
+        nn.Dropout(dropout),
+        nn.Linear(out_dim * 4, out_dim)
+    )
+
+
+class SpanQuery(nn.Module):
+
+    def __init__(self, hidden_size, max_width, trainable=True):
+        super().__init__()
+
+        self.query_seg = nn.Parameter(torch.randn(hidden_size, max_width))
+
+        nn.init.uniform_(self.query_seg, a=-1, b=1)
+
+        if not trainable:
+            self.query_seg.requires_grad = False
+
+        self.project = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.ReLU()
+        )
+
+    def forward(self, h, *args):
+        # h of shape [B, L, D]
+        # query_seg of shape [D, max_width]
+
+        span_rep = torch.einsum('bld, ds->blsd', h, self.query_seg)
+
+        return self.project(span_rep)
+
+
+class SpanMLP(nn.Module):
+
+    def __init__(self, hidden_size, max_width):
+        super().__init__()
+
+        self.mlp = nn.Linear(hidden_size, hidden_size * max_width)
+
+    def forward(self, h, *args):
+        # h of shape [B, L, D]
+        # query_seg of shape [D, max_width]
+
+        B, L, D = h.size()
+
+        span_rep = self.mlp(h)
+
+        span_rep = span_rep.view(B, L, -1, D)
+
+        return span_rep.relu()
+
+
+class SpanCAT(nn.Module):
+
+    def __init__(self, hidden_size, max_width):
+        super().__init__()
+
+        self.max_width = max_width
+
+        self.query_seg = nn.Parameter(torch.randn(128, max_width))
+
+        self.project = nn.Sequential(
+            nn.Linear(hidden_size + 128, hidden_size),
+            nn.ReLU()
+        )
+
+    def forward(self, h, *args):
+        # h of shape [B, L, D]
+        # query_seg of shape [D, max_width]
+
+        B, L, D = h.size()
+
+        h = h.view(B, L, 1, D).repeat(1, 1, self.max_width, 1)
+
+        q = self.query_seg.view(1, 1, self.max_width, -1).repeat(B, L, 1, 1)
+
+        span_rep = torch.cat([h, q], dim=-1)
+
+        span_rep = self.project(span_rep)
+
+        return span_rep
+
+
+class SpanConvBlock(nn.Module):
+    def __init__(self, hidden_size, kernel_size, span_mode='conv_normal'):
+        super().__init__()
+
+        if span_mode == 'conv_conv':
+            self.conv = nn.Conv1d(hidden_size, hidden_size,
+                                  kernel_size=kernel_size)
+
+            # initialize the weights
+            nn.init.kaiming_uniform_(self.conv.weight, nonlinearity='relu')
+
+        elif span_mode == 'conv_max':
+            self.conv = nn.MaxPool1d(kernel_size=kernel_size, stride=1)
+        elif span_mode == 'conv_mean' or span_mode == 'conv_sum':
+            self.conv = nn.AvgPool1d(kernel_size=kernel_size, stride=1)
+
+        self.span_mode = span_mode
+
+        self.pad = kernel_size - 1
+
+    def forward(self, x):
+
+        x = torch.einsum('bld->bdl', x)
+
+        if self.pad > 0:
+            x = F.pad(x, (0, self.pad), "constant", 0)
+
+        x = self.conv(x)
+
+        if self.span_mode == "conv_sum":
+            x = x * (self.pad + 1)
+
+        return torch.einsum('bdl->bld', x)
+
+
+class SpanConv(nn.Module):
+    def __init__(self, hidden_size, max_width, span_mode):
+        super().__init__()
+
+        kernels = [i + 2 for i in range(max_width - 1)]
+
+        self.convs = nn.ModuleList()
+
+        for kernel in kernels:
+            self.convs.append(SpanConvBlock(hidden_size, kernel, span_mode))
+
+        self.project = nn.Sequential(
+            nn.ReLU(),
+            nn.Linear(hidden_size, hidden_size)
+        )
+
+    def forward(self, x, *args):
+
+        span_reps = [x]
+
+        for conv in self.convs:
+            h = conv(x)
+            span_reps.append(h)
+
+        span_reps = torch.stack(span_reps, dim=-2)
+
+        return self.project(span_reps)
+
+
+class SpanEndpointsBlock(nn.Module):
+    def __init__(self, kernel_size):
+        super().__init__()
+
+        self.kernel_size = kernel_size
+
+    def forward(self, x):
+        B, L, D = x.size()
+
+        span_idx = torch.LongTensor(
+            [[i, i + self.kernel_size - 1] for i in range(L)]).to(x.device)
+
+        x = F.pad(x, (0, 0, 0, self.kernel_size - 1), "constant", 0)
+
+        # endrep
+        start_end_rep = torch.index_select(x, dim=1, index=span_idx.view(-1))
+
+        start_end_rep = start_end_rep.view(B, L, 2, D)
+
+        return start_end_rep
+
+
+class ConvShare(nn.Module):
+    def __init__(self, hidden_size, max_width):
+        super().__init__()
+
+        self.max_width = max_width
+
+        self.conv_weigth = nn.Parameter(
+            torch.randn(hidden_size, hidden_size, max_width))
+
+        nn.init.kaiming_uniform_(self.conv_weigth, nonlinearity='relu')
+
+        self.project = nn.Sequential(
+            nn.ReLU(),
+            nn.Linear(hidden_size, hidden_size)
+        )
+
+    def forward(self, x, *args):
+        span_reps = []
+
+        x = torch.einsum('bld->bdl', x)
+
+        for i in range(self.max_width):
+            pad = i
+            x_i = F.pad(x, (0, pad), "constant", 0)
+            conv_w = self.conv_weigth[:, :, :i + 1]
+            out_i = F.conv1d(x_i, conv_w)
+            span_reps.append(out_i.transpose(-1, -2))
+
+        out = torch.stack(span_reps, dim=-2)
+
+        return self.project(out)
+
+
+def extract_elements(sequence, indices):
+    B, L, D = sequence.shape
+    K = indices.shape[1]
+
+    # Expand indices to [B, K, D]
+    expanded_indices = indices.unsqueeze(2).expand(-1, -1, D)
+
+    # Gather the elements
+    extracted_elements = torch.gather(sequence, 1, expanded_indices)
+
+    return extracted_elements
+
+
+class SpanMarker(nn.Module):
+
+    def __init__(self, hidden_size, max_width, dropout=0.4):
+        super().__init__()
+
+        self.max_width = max_width
+
+        self.project_start = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 2, bias=True),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size * 2, hidden_size, bias=True),
+        )
+
+        self.project_end = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 2, bias=True),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size * 2, hidden_size, bias=True),
+        )
+
+        self.out_project = nn.Linear(hidden_size * 2, hidden_size, bias=True)
+
+    def forward(self, h, span_idx):
+        # h of shape [B, L, D]
+        # query_seg of shape [D, max_width]
+
+        B, L, D = h.size()
+
+        # project start and end
+        start_rep = self.project_start(h)
+        end_rep = self.project_end(h)
+
+        start_span_rep = extract_elements(start_rep, span_idx[:, :, 0])
+        end_span_rep = extract_elements(end_rep, span_idx[:, :, 1])
+
+        # concat start and end
+        cat = torch.cat([start_span_rep, end_span_rep], dim=-1).relu()
+
+        # project
+        cat = self.out_project(cat)
+
+        # reshape
+        return cat.view(B, L, self.max_width, D)
+
+
+class SpanMarkerV0(nn.Module):
+    """
+    Marks and projects span endpoints using an MLP.
+    """
+
+    def __init__(self, hidden_size: int, max_width: int, dropout: float = 0.4):
+        super().__init__()
+        self.max_width = max_width
+        self.project_start = create_projection_layer(hidden_size, dropout)
+        self.project_end = create_projection_layer(hidden_size, dropout)
+
+        self.out_project = create_projection_layer(hidden_size * 2, dropout, hidden_size)
+
+    def forward(self, h: torch.Tensor, span_idx: torch.Tensor) -> torch.Tensor:
+        B, L, D = h.size()
+
+        start_rep = self.project_start(h)
+        end_rep = self.project_end(h)
+
+        start_span_rep = extract_elements(start_rep, span_idx[:, :, 0])
+        end_span_rep = extract_elements(end_rep, span_idx[:, :, 1])
+
+        cat = torch.cat([start_span_rep, end_span_rep], dim=-1).relu()
+
+        return self.out_project(cat).view(B, L, self.max_width, D)
+
+
+class ConvShareV2(nn.Module):
+    def __init__(self, hidden_size, max_width):
+        super().__init__()
+
+        self.max_width = max_width
+
+        self.conv_weigth = nn.Parameter(
+            torch.randn(hidden_size, hidden_size, max_width)
+        )
+
+        nn.init.xavier_normal_(self.conv_weigth)
+
+    def forward(self, x, *args):
+        span_reps = []
+
+        x = torch.einsum('bld->bdl', x)
+
+        for i in range(self.max_width):
+            pad = i
+            x_i = F.pad(x, (0, pad), "constant", 0)
+            conv_w = self.conv_weigth[:, :, :i + 1]
+            out_i = F.conv1d(x_i, conv_w)
+            span_reps.append(out_i.transpose(-1, -2))
+
+        out = torch.stack(span_reps, dim=-2)
+
+        return out
+
+
+class SpanRepLayer(nn.Module):
+    """
+    Various span representation approaches
+    """
+
+    def __init__(self, hidden_size, max_width, span_mode, **kwargs):
+        super().__init__()
+
+        if span_mode == 'marker':
+            self.span_rep_layer = SpanMarker(hidden_size, max_width, **kwargs)
+        elif span_mode == 'markerV0':
+            self.span_rep_layer = SpanMarkerV0(hidden_size, max_width, **kwargs)
+        elif span_mode == 'query':
+            self.span_rep_layer = SpanQuery(
+                hidden_size, max_width, trainable=True)
+        elif span_mode == 'mlp':
+            self.span_rep_layer = SpanMLP(hidden_size, max_width)
+        elif span_mode == 'cat':
+            self.span_rep_layer = SpanCAT(hidden_size, max_width)
+        elif span_mode == 'conv_conv':
+            self.span_rep_layer = SpanConv(
+                hidden_size, max_width, span_mode='conv_conv')
+        elif span_mode == 'conv_max':
+            self.span_rep_layer = SpanConv(
+                hidden_size, max_width, span_mode='conv_max')
+        elif span_mode == 'conv_mean':
+            self.span_rep_layer = SpanConv(
+                hidden_size, max_width, span_mode='conv_mean')
+        elif span_mode == 'conv_sum':
+            self.span_rep_layer = SpanConv(
+                hidden_size, max_width, span_mode='conv_sum')
+        elif span_mode == 'conv_share':
+            self.span_rep_layer = ConvShare(hidden_size, max_width)
+        else:
+            raise ValueError(f'Unknown span mode {span_mode}')
+
+    def forward(self, x, *args):
+
+        return self.span_rep_layer(x, *args)

backup/modules/token_rep.py

@@ -0,0 +1,54 @@
+from typing import List
+
+import torch
+from flair.data import Sentence
+from flair.embeddings import TransformerWordEmbeddings
+from torch import nn
+from torch.nn.utils.rnn import pad_sequence
+
+
+# flair.cache_root = '/gpfswork/rech/pds/upa43yu/.cache'
+
+
+class TokenRepLayer(nn.Module):
+    def __init__(self, model_name: str = "bert-base-cased", fine_tune: bool = True, subtoken_pooling: str = "first",
+                 hidden_size: int = 768,
+                 add_tokens=["[SEP]", "[ENT]"]
+                 ):
+        super().__init__()
+
+        self.bert_layer = TransformerWordEmbeddings(
+            model_name,
+            fine_tune=fine_tune,
+            subtoken_pooling=subtoken_pooling,
+            allow_long_sentences=True
+        )
+
+        # add tokens to vocabulary
+        self.bert_layer.tokenizer.add_tokens(add_tokens)
+
+        # resize token embeddings
+        self.bert_layer.model.resize_token_embeddings(len(self.bert_layer.tokenizer))
+
+        bert_hidden_size = self.bert_layer.embedding_length
+
+        if hidden_size != bert_hidden_size:
+            self.projection = nn.Linear(bert_hidden_size, hidden_size)
+
+    def forward(self, tokens: List[List[str]], lengths: torch.Tensor):
+        token_embeddings = self.compute_word_embedding(tokens)
+
+        if hasattr(self, "projection"):
+            token_embeddings = self.projection(token_embeddings)
+
+        B = len(lengths)
+        max_length = lengths.max()
+        mask = (torch.arange(max_length).view(1, -1).repeat(B, 1) < lengths.cpu().unsqueeze(1)).to(
+            token_embeddings.device).long()
+        return {"embeddings": token_embeddings, "mask": mask}
+
+    def compute_word_embedding(self, tokens):
+        sentences = [Sentence(i) for i in tokens]
+        self.bert_layer.embed(sentences)
+        token_embeddings = pad_sequence([torch.stack([t.embedding for t in k]) for k in sentences], batch_first=True)
+        return token_embeddings