Update README.md

README.md

@@ -18,21 +18,33 @@ Provence is a lightweight **context pruning model** for retrieval-augmented gene
 * *Backbone model*: [DeBERTav3-reranker](https://huggingface.co/naver/trecdl22-crossencoder-debertav3) (trained from [DeBERTa-v3-large](https://huggingface.co/microsoft/deberta-v3-large))  
 * *Model size*: 430 million parameters  
 * *Context length*: 512 tokens  
-* *Other model variants*: *TODO* 
 
 ## Usage
 
+Pruning a single context for a single question:
+
 ```python
 from transformers import AutoModel
 
 provence = AutoModel.from_pretrained("naver/provence-reranker-debertav3-v1", trust_remote_code=True)
 
-context = [["Shepherd’s pie. History. In early cookery books, the dish was a means of using leftover roasted meat of any kind, and the pie dish was lined on the sides and bottom with mashed potato, as well as having a mashed potato crust on top. Variations and similar dishes. Other potato-topped pies include: The modern ”Cumberland pie” is a version with either beef or lamb and a layer of bread- crumbs and cheese on top. In medieval times, and modern-day Cumbria, the pastry crust had a filling of meat with fruits and spices.. In Quebec, a varia- tion on the cottage pie is called ”Paˆte ́ chinois”. It is made with ground beef on the bottom layer, canned corn in the middle, and mashed potato on top.. The ”shepherdess pie” is a vegetarian version made without meat, or a vegan version made without meat and dairy.. In the Netherlands, a very similar dish called ”philosopher’s stew” () often adds ingredients like beans, apples, prunes, or apple sauce.. In Brazil, a dish called in refers to the fact that a manioc puree hides a layer of sun-dried meat."]]
-query = ['what goes on the bottom of shepherd’s pie']
+context = "Shepherd’s pie. History. In early cookery books, the dish was a means of using leftover roasted meat of any kind, and the pie dish was lined on the sides and bottom with mashed potato, as well as having a mashed potato crust on top. Variations and similar dishes. Other potato-topped pies include: The modern ”Cumberland pie” is a version with either beef or lamb and a layer of bread- crumbs and cheese on top. In medieval times, and modern-day Cumbria, the pastry crust had a filling of meat with fruits and spices.. In Quebec, a varia- tion on the cottage pie is called ”Paˆte ́ chinois”. It is made with ground beef on the bottom layer, canned corn in the middle, and mashed potato on top.. The ”shepherdess pie” is a vegetarian version made without meat, or a vegan version made without meat and dairy.. In the Netherlands, a very similar dish called ”philosopher’s stew” () often adds ingredients like beans, apples, prunes, or apple sauce.. In Brazil, a dish called in refers to the fact that a manioc puree hides a layer of sun-dried meat."
+question = 'What goes on the bottom of Shepherd’s pie?'
+
+provence_output = provence.process(question, context)
+# print(f"Provence Output: {provence_output}")
+# Provence Output: {'reranking_score': 3.022725, pruned_context': 'In early cookery books, the dish was a means of using leftover roasted meat of any kind, and the pie dish was lined on the sides and bottom with mashed potato, as well as having a mashed potato crust on top.']]
+```
+
+You can also pass a list of questions and a list of lists of contexts (multiple contexts per question to be pruned) for batched processing.
+
+Setting `always_select_title=True` will keep the first sentence "Shepherd’s pie". This is especially useful for Wikipedia articles where the title is often needed to understand the context.
+More details on how the title is defined are given below.
 
-pruned_context = provence.process(context, query)
-# print(f"Pruned context: {pruned_context}")
-# Pruned context: [['Shepherd’s pie. In early cookery books, the dish was a means of using leftover roasted meat of any kind, and the pie dish was lined on the sides and bottom with mashed potato, as well as having a mashed potato crust on top.']]
+```python
+provence_output = provence.process(question, context, always_select_title=True)
+# print(f"Provence Output: {provence_output}")
+# Provence Output: {'reranking_score': 3.022725, pruned_context': 'Shepherd’s pie. In early cookery books, the dish was a means of using leftover roasted meat of any kind, and the pie dish was lined on the sides and bottom with mashed potato, as well as having a mashed potato crust on top.']]
 ```
 
 Training code, as well as RAG experiments with Provence can be found in the [BERGEN](https://github.com/naver/bergen) library. 
@@ -40,11 +52,11 @@ Training code, as well as RAG experiments with Provence can be found in the [BER
 ## Model interface
 
 Interface of the `process` function:
-* `questions`: `List[str]`: a list of input questions
-* `contexts`: `List[List[str]]`: a list of retrieved contexts, provided in a list for each question. `len(contexts)` should be equal to `len(questions)`
-*  `titles`: `Optional[Union[List[List[str]], str]]`, _default: “first_sentence”_: an optional list of titles for retrieved contexts, same shape as `contexts`. If it is equal to `first_sentence`, then the first sentence of each context is assumed to be the title. If None, then it is assumed that no titles are provided. Titles are only used if `always_select_title=True`.
+* `question`: `Union[List[str], str]`: an input question (str) or a list of input questions (for batched processing)
+* `context`: `Union[List[List[str]], [List[str]], str]`: context(s) to be pruned. This can be either a single string (in case of a singe str question), or a list of lists contexts (a list of contexts per question), with `len(contexts)` equal to `len(questions)`
+*  `title`: `Optional[Union[List[List[str]], str]]`, _default: “first_sentence”_: an optional argument for defining titles. If `title=first_sentence`, then the first sentence of each context is assumed to be the title. If `title=None`, then it is assumed that no titles are provided. Titles can be also passed as a list of lists of str, i.e. titles shaped the same way as contexts. Titles are only used if `always_select_title=True`.
 * `threshold` _(float, $ \in [0, 1]$, default 0.1)_: which threshold to use for context pruning. We recommend 0.1 for more conservative pruning (no performance drop or lowest performance drops) and 0.5 for higher compression, but this value can be further tuned to meet the specific use case requirements.
-* `always_select_title` _(bool, default: True)_: if True, the first sentence (title) will always be selected. This is important, e.g., for Wikipedia passages, to provide proper context for the next sentences.
+* `always_select_title` _(bool, default: True)_: if True, the first sentence (title) will always be selected. This is important, e.g., for Wikipedia passages, to provide proper contextualization for the next sentences.
 * `batch_size` (int, default: 32)
 * `reorder` _(bool, default: False)_: if True, the provided contexts for each question will be reordered according to the computed question-passage relevance scores. If False, the original user-provided order of contexts will be preserved.
 * `top_k` _(int, default: 5)_: if `reorder=True`, specifies the number of top-ranked passages to keep for each question.