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