Upload modeling file

added_tokens.json

@@ -0,0 +1,3 @@
+{
+  "[MASK]": 128000
+}

config.json

@@ -0,0 +1 @@
+{"_name_or_path": "provence", "architectures": ["Provence"], "auto_map": {"AutoConfig": "modeling_provence.ProvenceConfig", "AutoModel": "modeling_provence.Provence"}, "attention_probs_dropout_prob": 0.1, "hidden_act": "gelu", "hidden_dropout_prob": 0.1, "hidden_size": 1024, "id2label": {"0": "LABEL_0"}, "initializer_range": 0.02, "intermediate_size": 4096, "label2id": {"LABEL_0": 0}, "layer_norm_eps": 1e-07, "max_position_embeddings": 512, "max_relative_positions": -1, "model_type": "Provence", "norm_rel_ebd": "layer_norm", "num_attention_heads": 16, "num_hidden_layers": 24, "pad_token_id": 0, "pooler_dropout": 0, "pooler_hidden_act": "gelu", "pooler_hidden_size": 1024, "pos_att_type": ["p2c", "c2p"], "position_biased_input": false, "position_buckets": 256, "relative_attention": true, "share_att_key": true, "torch_dtype": "float32", "transformers_version": "4.45.1", "type_vocab_size": 0, "vocab_size": 128100}

modeling_provence.py

@@ -0,0 +1,456 @@
+import time
+import string
+from typing import Optional, Union, Tuple, List
+from dataclasses import dataclass
+from tqdm import tqdm
+import warnings
+import nltk
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from torch.nn.utils.rnn import pad_sequencea
+from transformers import AutoTokenizer
+from transformers import DebertaV2PreTrainedModel, DebertaV2Model, PretrainedConfig
+from transformers.models.deberta_v2.modeling_deberta_v2 import (
+    StableDropout,
+    ContextPooler,
+)
+from transformers.modeling_outputs import ModelOutput
+
+
+@dataclass
+class RankingCompressionOutput(ModelOutput):
+
+    compression_logits: torch.FloatTensor = None
+    ranking_scores: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+"""adapted from https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L1357
+"""
+
+
+class ProvenceConfig(PretrainedConfig):
+
+    model_type = "Provence"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+class Provence(DebertaV2PreTrainedModel):
+
+    config_class = ProvenceConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        num_labels = getattr(config, "num_labels", 2)
+        self.num_labels = num_labels
+        self.deberta = DebertaV2Model(config)
+        self.pooler = ContextPooler(config)
+        output_dim = self.pooler.output_dim
+
+        ### RANKING LAYER
+        self.classifier = nn.Linear(output_dim, num_labels)
+        drop_out = getattr(config, "cls_dropout", None)
+        drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
+        self.dropout = StableDropout(drop_out)
+
+        ### COMPRESSION LAYER: another head
+        token_dropout = drop_out
+        self.token_dropout = nn.Dropout(token_dropout)
+        self.token_classifier = nn.Linear(
+            config.hidden_size, 2
+        )  # => hard coded number of labels
+        self.name = "Provence"
+        self.tokenizer = AutoTokenizer.from_pretrained(config._name_or_path)
+        self.max_len = config.max_position_embeddings
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> RankingCompressionOutput:
+        outputs = self.deberta(
+            input_ids,
+            attention_mask=attention_mask,
+        )
+
+        encoder_layer = outputs[0]
+        pooled_output = self.pooler(encoder_layer)
+        pooled_output = self.dropout(pooled_output)
+        ranking_logits = self.classifier(pooled_output)
+        compression_logits = self.token_classifier(self.token_dropout(encoder_layer))
+        ranking_scores = ranking_logits[
+            :, 0
+        ].squeeze()  # select first dim of logits for ranking scores
+
+        return RankingCompressionOutput(
+            compression_logits=compression_logits,
+            ranking_scores=ranking_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def process(
+        self,
+        contexts: List[List[str]],
+        queries: List[List[str]],
+        titles: Optional[Union[List[str], str]] = "first_sentence",
+        batch_size=32,
+        threshold=0.01,
+        always_select_title=True,
+        reorder=False,
+        top_k=5,
+        enable_warnings=True,
+    ):
+
+        assert (
+            titles == "first_sentence"
+            or titles == None
+            or type(titles) == list
+            and len(titles) == len(queries)
+        ), "Variable 'titles' must be 'first_sentence' or a list of strings of the same length as 'queries'"
+        if type(titles) == list:
+            assert all(
+                [
+                    len(titles_item) == len(contexts_item)
+                    for titles_item, contexts_item in zip(contexts, titles)
+                ]
+            ), "Each list in 'titles' must have the same length as the corresponding list in 'context'"
+        assert len(queries) == len(
+            contexts
+        ), "Lists 'queries' and 'contexts' must have same lengths"
+        times = []
+        t0 = time.time()
+        dataset = TestDataset(
+            queries=queries,
+            contexts=contexts,
+            titles=titles,
+            tokenizer=self.tokenizer,
+            max_len=self.max_len,
+            enable_warnings=enable_warnings,
+        )
+        times.append(["testdataset", time.time() - t0])
+        t0 = time.time()
+        selected_contexts = [
+            [{0: contexts[i][j]} for j in range(len(contexts[i]))]
+            for i in range(len(queries))
+        ]
+        reranking_scores = [
+            [None for j in range(len(contexts[i]))] for i in range(len(queries))
+        ]
+        times.append(["create arrays", time.time() - t0])
+        t0 = time.time()
+        with torch.no_grad():
+            for batch_start in tqdm(
+                range(0, len(dataset), batch_size), desc="Pruning contexts..."
+            ):
+                t1 = time.time()
+                qis = dataset.qis[batch_start : batch_start + batch_size]
+                cis = dataset.cis[batch_start : batch_start + batch_size]
+                sis = dataset.sis[batch_start : batch_start + batch_size]
+                sent_coords = dataset.sent_coords[
+                    batch_start : batch_start + batch_size
+                ]
+                ids_list = dataset.ids[batch_start : batch_start + batch_size]
+                ids = pad_sequence(
+                    ids_list, batch_first=True, padding_value=dataset.pad_idx
+                ).to(self.device)
+                mask = (ids != dataset.pad_idx).to(self.device)
+                times.append(["torch stack", time.time() - t1])
+                t1 = time.time()
+                outputs = self.forward(ids, mask)
+                scores = F.softmax(outputs["compression_logits"].cpu(), dim=-1)[:, :, 1]
+                token_preds = scores > threshold
+                reranking_scrs = (
+                    outputs["ranking_scores"].cpu().numpy()
+                )  # get first score
+                if len(reranking_scrs.shape) == 0:
+                    reranking_scrs = reranking_scrs[None]
+                times.append(["forward pass", time.time() - t1])
+                t1 = time.time()
+                for (
+                    ids_list_,
+                    token_preds_,
+                    rerank_score,
+                    qi,
+                    ci,
+                    si,
+                    sent_coords_,
+                ) in zip(
+                    ids_list, token_preds, reranking_scrs, qis, cis, sis, sent_coords
+                ):
+
+                    selected_mask = sentence_rounding(
+                        token_preds_.cpu().numpy(),
+                        np.array(sent_coords_),
+                        threshold=threshold,
+                        always_select_title=always_select_title
+                        and si == 0
+                        and titles != None,
+                    )
+                    assert len(selected_mask) == len(token_preds_)
+                    selected_contexts[qi][ci][si] = ids_list_[
+                        selected_mask[: len(ids_list_)]
+                    ]
+                    if si == 0:
+                        reranking_scores[qi][ci] = rerank_score
+                times.append(["postprocessing", time.time() - t1])
+        t1 = time.time()
+        for i in range(len(queries)):
+            for j in range(len(contexts[i])):
+                if type(selected_contexts[i][j][0]) != str:
+                    toks = torch.cat(
+                        [
+                            ids_
+                            for _, ids_ in sorted(
+                                selected_contexts[i][j].items(), key=lambda x: x[0]
+                            )
+                        ]
+                    )
+                    selected_contexts[i][j] = self.tokenizer.decode(
+                        toks,
+                        skip_special_tokens=True,
+                        clean_up_tokenization_spaces=False,
+                    )
+                else:
+                    selected_contexts[i][j] = selected_contexts[i][j][0]
+            if reorder:
+                print(reranking_scores[qi])
+                print(np.sort(reranking_scores[i])[::-1][:top_k])
+                idxs = np.argsort(reranking_scores[i])[::-1][:top_k]
+                selected_contexts[i] = [selected_contexts[i][j] for j in idxs]
+                reranking_scores[i] = [reranking_scores[i][j] for j in idxs]
+        times.append(["postpostprocessing", time.time() - t1])
+        times.append(["total inference", time.time() - t0])
+        return selected_contexts  # , reranking_scores, times
+
+
+# Some utils functions
+
+
+def sentence_rounding(predictions, chunks, threshold, always_select_title=True):
+    """
+    predictions: a binary vector containing 1 for tokens which were selected and 0s otherwise
+    chunks: a list of pairs [start, end] of sentence, i.e. sentence is in coordinates predictions[start:end]
+    the functions
+    """
+    cumulative_sum = np.cumsum(predictions)
+    chunk_sums = cumulative_sum[chunks[:, 1] - 1] - np.where(
+        chunks[:, 0] > 0, cumulative_sum[chunks[:, 0] - 1], 0
+    )
+    chunk_lengths = chunks[:, 1] - chunks[:, 0]
+    chunk_means = chunk_sums / chunk_lengths
+    if always_select_title:
+        chunk_means[0] = 1
+    means = np.hstack((np.zeros(1), chunk_means, np.zeros(1)))
+    repeats = np.hstack(
+        ([chunks[0][0]], chunk_lengths, [predictions.shape[0] - chunks[-1][1]])
+    )
+    return np.repeat(means, repeats) > threshold
+
+
+def normalize(s: str) -> str:
+    def white_space_fix(text):
+        return " ".join(text.split())
+
+    def remove_punc(text):
+        exclude = set(string.punctuation)
+        return "".join(ch for ch in text if ch not in exclude)
+
+    def lower(text):
+        return text.lower()
+
+    return white_space_fix(remove_punc(lower(s)))
+
+
+def sent_split_and_tokenize(text, tokenizer, max_len):
+    sents_nltk = nltk.sent_tokenize(text)
+    sents = []
+    for j, sent_nltk in enumerate(sents_nltk):
+        tokinput = (" " if j != 0 else "") + sent_nltk
+        tok = tokenizer.encode(tokinput, add_special_tokens=False)
+        ltok = len(tok)
+        if ltok == 0:
+            continue
+        if ltok <= max_len:
+            sents.append(tok)
+        else:
+            for begin in range(0, ltok, max_len):
+                sents.append(tok[begin : begin + max_len])
+    return sents
+
+
+class TestDataset(Dataset):
+    def __init__(
+        self,
+        queries,
+        contexts,
+        tokenizer,
+        max_len=512,
+        titles="first_sentence",
+        enable_warnings=True,
+    ):
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+        self.pad_idx = 0
+        self.cls_idx = [1]
+        self.sep_idx = [2]
+        self.eos = [2]
+        # hardcoded deberta-specific indexes
+        self.nb_spe_tok = len(self.cls_idx) + len(self.sep_idx)
+        self.enable_warnings = enable_warnings
+        self.unusual_query_length = (
+            self.max_len // 2
+        )  # TODO: change to data-driven value
+        self.unusual_title_len = self.max_len // 2  # TODO: change to data-driven value
+        self.create_dataset(contexts, queries, titles)
+        self.len = len(self.cis)
+
+    def create_dataset(self, contexts, queries, titles="first_sentence"):
+        self.qis = []
+        self.cis = []
+        self.sis = []
+        self.sent_coords = []
+        self.cntx_coords = []
+        self.ids = []
+        if self.enable_warnings:
+            warnings_dict = {
+                "zero_len_query": set(),
+                "too_long_query": set(),
+                "unusually_long_query": set(),
+                "unusually_long_title": set(),
+                "split_context": set(),
+            }
+        for i, query in enumerate(queries):
+            tokenized_query = self.tokenizer.encode(
+                normalize(query), add_special_tokens=False
+            )
+            # normalize query because all training data has normalized queries
+            query_len = len(tokenized_query)
+            if query_len == 0:
+                if self.enable_warnings:
+                    warnings_dict["zero_len_query"].add(i)
+                continue
+            elif query_len >= self.max_len - self.nb_spe_tok - 1:  # -1 for eos
+                if self.enable_warnings:
+                    warnings_dict["too_long_query"].add(i)
+                continue
+            elif query_len >= self.unusual_query_length:
+                if self.enable_warnings:
+                    warnings_dict["unusually_long_query"].add(i)
+            left_0 = len(tokenized_query) + self.nb_spe_tok
+            tokenized_seq_0 = self.cls_idx + tokenized_query + self.sep_idx
+            max_len = self.max_len - left_0 - 1
+            for j, cntx in enumerate(contexts[i]):
+                title = titles[i][j] if type(titles) == list else titles
+                tokenized_sents = sent_split_and_tokenize(cntx, self.tokenizer, max_len)
+                # each (sent + query + special tokens) <= max_len
+                if title is not None and title != "first_sentence":
+                    tokenized_title = self.tokenizer.encode(
+                        title, add_special_tokens=False
+                    )
+                    ltok = len(tokenized_title)
+                    if ltok == 0:
+                        pass
+                    elif ltok <= max_len:
+                        tokenized_sents = [tokenized_title] + tokenized_sents
+                    else:
+                        if self.enable_warnings and ltok >= self.unusual_title_len:
+                            warnings_dict["unusually_long_title"].add(i)
+                        tokenized_sents = [
+                            tokenized_title[begin : begin + max_len]
+                            for begin in range(0, ltok, max_len)
+                        ] + tokenized_sents
+                tokenized_seq = tokenized_seq_0
+                left = left_0
+                sent_coords = []
+                block = 0
+                for idx, tokenized_sent in enumerate(tokenized_sents):
+                    l = len(tokenized_sent)
+                    if left + l <= self.max_len - 1:
+                        sent_coords.append([left, left + l])
+                        tokenized_seq = tokenized_seq + tokenized_sent
+                        left += l
+                    else:
+                        if self.enable_warnings:
+                            warnings_dict["split_context"].add(i)
+                        if len(tokenized_seq) > left_0:
+                            tokenized_seq = tokenized_seq + self.eos
+                            self.qis.append(i)
+                            self.cis.append(j)
+                            self.sis.append(block)
+                            self.sent_coords.append(sent_coords)
+                            self.cntx_coords.append(
+                                [sent_coords[0][0], sent_coords[-1][1]]
+                            )
+                            self.ids.append(torch.tensor(tokenized_seq))
+                        tokenized_seq = tokenized_seq_0 + tokenized_sent
+                        sent_coords = [[left_0, left_0 + l]]
+                        left = left_0 + l
+                        block += 1
+                if len(tokenized_seq) > left_0:
+                    tokenized_seq = tokenized_seq + self.eos
+                    self.qis.append(i)
+                    self.cis.append(j)
+                    self.sis.append(block)
+                    self.sent_coords.append(sent_coords)
+                    self.cntx_coords.append([sent_coords[0][0], sent_coords[-1][1]])
+                    self.ids.append(torch.tensor(tokenized_seq))
+        if self.enable_warnings:
+            self.print_warnings(warnings_dict, len(queries))
+
+    def __len__(self):
+        return len(self.ids)
+
+    def print_warnings(self, warnings_dict, N):
+        n = len(warnings_dict["zero_len_query"])
+        info = " You can suppress Provence warnings by setting enable_warnings=False."
+        if n > 0:
+            ex = list(warnings_dict["zero_len_query"])[:10]
+            warnings.warn(
+                f"{n} out of {N} queries have zero length, e.g. at indexes {ex}. "
+                "These examples will be skipped in context pruning, "
+                "their contexts will be kept as is." + info
+            )
+        n = len(warnings_dict["too_long_query"])
+        if n > 0:
+            ex = list(warnings_dict["too_long_query"])[:10]
+            warnings.warn(
+                f"{n} out of {N} queries are too long for context length {self.max_len}, "
+                f"e.g. at indexes {ex}. These examples will be skipped in context pruning, "
+                "their contexts will be kept as is." + info
+            )
+        n = len(warnings_dict["unusually_long_query"])
+        if n > 0:
+            ex = list(warnings_dict["unusually_long_query"])[:10]
+            warnings.warn(
+                f"{n} out of {N} queries are longer than {self.unusual_query_length} tokens, "
+                f"e.g. at indexes {ex}. These examples will processed as usual in context pruning, "
+                "but the quality of context pruning could be reduced." + info
+            )
+        n = len(warnings_dict["unusually_long_title"])
+        if n > 0:
+            ex = list(warnings_dict["unusually_long_title"])[:10]
+            warnings.warn(
+                f"{n} out of {N} titles are longer than {self.unusual_title_length} tokens, "
+                f"e.g. at indexes {ex}. These examples will processed as usual in context pruning, "
+                "but the quality of context pruning could be reduced." + info
+            )
+        n = len(warnings_dict["split_context"])
+        if n > 0:
+            ex = list(warnings_dict["split_context"])[:10]
+            warnings.warn(
+                f"{n} out of {N} contexts were split into several pieces for context pruning, "
+                f"due to a limited context length of Provence which is equal to {self.max_len}. "
+                "This could potentially reduce the quality of context pruning. "
+                "You could consider checking and reducing lengths of contexts, queries, or titles."
+                + info
+            )

special_tokens_map.json

@@ -0,0 +1,15 @@
+{
+  "bos_token": "[CLS]",
+  "cls_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,58 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128000": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "[CLS]",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "do_lower_case": false,
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "sp_model_kwargs": {},
+  "split_by_punct": false,
+  "tokenizer_class": "DebertaV2Tokenizer",
+  "unk_token": "[UNK]",
+  "vocab_type": "spm"
+}