Add semantic highlighter model files with LFS for large files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model.safetensors filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,157 @@
----
-license: apache-2.0
----
+# opensearch-semantic-highlighter
+
+## Overview
+
+The OpenSearch semantic highlighter is a trained classifier that takes a document and query as input and returns a binary score for each sentence in the document indicating its relevance to the query. 
+
+## Usage
+
+The model is intended to be used within the OpenSearch cluster. However, for illustrative purposes, we include an example of how it can be used outside of OpenSearch:
+
+```python
+import nltk
+import torch
+import numpy as np
+from datasets import Dataset
+from functools import partial
+from torch.utils.data import DataLoader
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Union
+from torch.nn.utils.rnn import pad_sequence
+from transformers import AutoTokenizer
+from highlighter_model_tracing import TraceableBertTaggerForSentenceExtractionWithBackoff
+
+@dataclass
+class DataCollatorWithPadding:
+    pad_kvs: Dict[str, Union[int, float]] = field(default_factory=dict)
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
+        first = features[0]
+        batch = {}
+
+        # pad and collate keys in self.pad_kvs
+        for key, pad_value in self.pad_kvs.items():
+            if key in first and first[key] is not None:
+                batch[key] = pad_sequence(
+                    [torch.tensor(f[key]) for f in features],
+                    batch_first=True,
+                    padding_value=pad_value,
+                )
+
+        # collate remaining keys assuming that the values can be stacked
+        for k, v in first.items():
+            if k not in self.pad_kvs and v is not None and isinstance(v, torch.Tensor):
+                batch[k] = torch.stack([f[k] for f in features])
+
+        return batch
+
+
+def prepare_input_features(
+    tokenizer, examples, max_seq_length=512, stride=128, padding=False
+):
+
+    # jointly tokenize questions and context
+    tokenized_examples = tokenizer(
+        examples["question"],
+        examples["context"],
+        truncation="only_second",
+        max_length=max_seq_length,
+        stride=stride,
+        return_overflowing_tokens=True,
+        padding=padding,
+        is_split_into_words=True,
+    )
+
+    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
+    tokenized_examples["example_id"] = []
+    tokenized_examples["word_ids"] = []
+    tokenized_examples["sentence_ids"] = []
+
+    # process model inputs
+    for i, sample_index in enumerate(sample_mapping):
+        word_ids = tokenized_examples.word_ids(i)
+        word_level_sentence_ids = examples["word_level_sentence_ids"][sample_index]
+
+        sequence_ids = tokenized_examples.sequence_ids(i)
+        token_start_index = 0
+        while sequence_ids[token_start_index] != 1:
+            token_start_index += 1
+
+        sentences_ids = [-100] * token_start_index
+        for word_idx in word_ids[token_start_index:]:
+            if word_idx is not None:
+                sentences_ids.append(word_level_sentence_ids[word_idx])
+            else:
+                sentences_ids.append(-100)
+
+        tokenized_examples["sentence_ids"].append(sentences_ids)
+        tokenized_examples["example_id"].append(examples["id"][sample_index])
+        tokenized_examples["word_ids"].append(word_ids)
+
+    return tokenized_examples
+
+
+# example highlighting case, from OpenSearch documentation
+query = "When does OpenSearch use text reanalysis for highlighting?"
+document = "To highlight the search terms, the highlighter needs the start and end character offsets of each term. The offsets mark the term’s position in the original text. The highlighter can obtain the offsets from the following sources: Postings: When documents are indexed, OpenSearch creates an inverted search index—a core data structure used to search for documents. Postings represent the inverted search index and store the mapping of each analyzed term to the list of documents in which it occurs. If you set the index_options parameter to offsets when mapping a text field, OpenSearch adds each term’s start and end character offsets to the inverted index. During highlighting, the highlighter reruns the original query directly on the postings to locate each term. Thus, storing offsets makes highlighting more efficient for large fields because it does not require reanalyzing the text. Storing term offsets requires additional disk space, but uses less disk space than storing term vectors. Text reanalysis: In the absence of both postings and term vectors, the highlighter reanalyzes text in order to highlight it. For every document and every field that needs highlighting, the highlighter creates a small in-memory index and reruns the original query through Lucene’s query execution planner to access low-level match information for the current document. Reanalyzing the text works well in most use cases. However, this method is more memory and time intensive for large fields."
+
+# sentence-level parsing
+sentence_ids = []
+context = []
+document_sents = nltk.sent_tokenize(document)
+for sent_id, sent in enumerate(document_sents):
+    sent_words = sent.split(' ')
+    context += sent_words
+    sentence_ids += [sent_id] * len(sent_words)
+
+# format example highlighting case as a Dataset
+example_dataset = Dataset.from_dict({'question': [[query]], 
+               'context': [context], 
+               'word_level_sentence_ids': [sentence_ids],
+               'id': [0]}) 
+
+# prepare to featurize the raw text data
+base_model_id = "bert-base-uncased"
+tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+collator = DataCollatorWithPadding(
+        pad_kvs={
+            "input_ids": 0,
+            "token_type_ids": 0,
+            "attention_mask": 0,
+            "sentence_ids": -100,
+            "sentence_labels": -100,
+        }
+    )
+preprocess_fn = partial(prepare_input_features, tokenizer)
+
+# featurize 
+example_dataset = example_dataset.map(
+    preprocess_fn,
+    batched=True,
+    remove_columns=example_dataset.column_names,
+    desc="Preparing model inputs",
+)
+loader = DataLoader(example_dataset, batch_size=1, collate_fn=collator)
+batch = next(iter(loader))
+
+# load model and get sentence highlights
+model = TraceableBertTaggerForSentenceExtractionWithBackoff.from_pretrained('opensearch-project/opensearch-semantic-highlighter-v1')
+highlights = model(
+    batch["input_ids"],
+    batch["attention_mask"],
+    batch["token_type_ids"],
+    batch["sentence_ids"],
+)
+highlighted_sentences = [document_sents[x] for x in highlights[0]]
+print(highlighted_sentences)
+```
+
+
+## License
+
+This project is licensed under the [Apache v2.0 License](https://github.com/opensearch-project/neural-search/blob/main/LICENSE).
+
+
+## Copyright
+
+Copyright OpenSearch Contributors. See [NOTICE](https://github.com/opensearch-project/neural-search/blob/main/NOTICE) for details.

config.json

@@ -0,0 +1,26 @@
+{
+  "_name_or_path": "bert-base-uncased",
+  "architectures": [
+    "BertTaggerForSentenceExtraction"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.2",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae65ebbec6562c9badc6ae4bd5be5d1b55536ec840658c1d2821d0c6258b7146
+size 437958648

tokenizer_config.json

@@ -0,0 +1 @@
+{"do_lower_case": true, "model_max_length": 512}