update

README.md

@@ -9,107 +9,7 @@ app_file: app.py
 pinned: false
 ---
 
-<style>
-.transformer {
-  display: inline-block;
-  background: #8facdb;
-  position: relative;
-  height: 60px;
-  line-height: 60px;
-  padding: 0 24px;
-  margin: 0 18px;
-  color: #333;
-  cursor: help;
-}
-.transformer::before {
-  content: "";
-  position: absolute;
-  bottom: 0;
-  top: 0;
-  left: -15px;
-  border-top: 30px solid #8facdb;
-  border-bottom: 30px solid #8facdb;
-  border-left: 15px solid transparent;
-}
-.transformer::after {
-  content: "";
-  position: absolute;
-  bottom: 0;
-  top: 0;
-  right: -15px;
-  border-top: 30px solid transparent;
-  border-bottom: 30px solid transparent;
-  border-left: 15px solid #8facdb;
-}
-.transformer.boring {
-  background: #ddd;
-}
-.transformer.boring::before {
-  border-top-color: #ddd;
-  border-bottom-color: #ddd;
-}
-.transformer.boring::after {
-  border-left-color: #ddd;
-}
-.df {
-  width: 24px;
-  line-height: 24px;
-  text-align: center;
-  border: 3px double #888;
-  background-color: #eee;
-  color: #333;
-  border-radius: 4px;
-  display: inline-block;
-  box-sizing: content-box;
-  cursor: help;
-  margin: 0 -25px;
-  opacity: 0.5;
-  z-index: 1;
-  position: relative;
-}
-.df:hover {
-  opacity: 1;
-}
-.pipeline {
-  text-align: center;
-}
-.artefact {
-  width: 32px;
-  line-height: 32px;
-  background: #eee;
-  display: inline-block;
-  box-sizing: content-box;
-  cursor: help;
-  margin: 0 -25px;
-  z-index: 1;
-  opacity: 0.5;
-  position: relative;
-  color: #333;
-  text-align: center;
-  border: 3px double #888;
-  border-radius: 50%
-}
-.artefact:hover {
-  opacity: 1;
-}
-.transformer .artefact {
-  bottom: -12px;
-  left: 50%;
-  margin-left: -16px;
-}
-</style>
+# 🐕 PyTerrier: SPLADE
 
 This is a demonstration of [PyTerrier's SPLADE package](https://github.com/cmacdonald/pyt_splade). The SPLADE model encodes queries and documents
 into sparse representations, which can then be used for indexing and retrieval.
-
-### Query Encoding
-
-Let's start by exploring SPLADE's query encoder. The query encoder is a `Q→Q` (query rewriting, query-to-query) transformer, and can be used in pipelines accordingly.
-It maps a query string into [MatchOp](https://terrier-core.readthedocs.io/en/latest/querylanguage.html#matching-op-query-language) query with terms from the
-query re-weighted and weighted expansion terms added.
-
-<div class="pipeline">
-  <div class="df" title="Query Frame">Q</div>
-  <div class="transformer" title="SPLADE Query Transformer">SPLADE</div>
-  <div class="df" title="Query Frame">Q</div>
-</div>

app.py

@@ -1,14 +1,13 @@
 import re
 import json
-import base64
 import pandas as pd
 import gradio as gr
 import pyterrier as pt
 pt.init()
 import pyt_splade
-factory = pyt_splade.SpladeFactory()
-pipe_queries = factory.query()
-pipe_docs = factory.indexing()
+from pyterrier_gradio import Demo, MarkdownFile, interface, df2code, code2md
+factory_max = pyt_splade.SpladeFactory(agg='max')
+factory_sum = pyt_splade.SpladeFactory(agg='sum')
 
 COLAB_NAME = 'pyterrier_splade.ipynb'
 COLAB_INSTALL = '''
@@ -16,30 +15,6 @@ COLAB_INSTALL = '''
 !pip install -q git+https://github.com/seanmacavaney/pyt_splade@misc
 '''.strip()
 
-def df2code(df):
-  rows = []
-  for row in df.itertuples(index=False):
-    rows.append(f'  {dict(row._asdict())},')
-  rows = '\n'.join(rows)
-  return f'''pd.DataFrame([
-{rows}
-])'''
-
-def code2colab(code):
-  enc_code = base64.b64encode((COLAB_INSTALL + '\n\n' + code.strip()).encode()).decode()
-  dec = base64.b64decode(enc_code)
-  url = f'https://colaburl.macavaney.us/?py64={enc_code}&name={COLAB_NAME}'
-  return f'<div style="text-align: center; margin-bottom: -16px;"><a href="{url}" rel="nofollow" target="_blank"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab" style="margin: 0; display: inline-block;" /></a></div>'
-
-def code2md(code):
-  return f'''
-{code2colab(code)}
-
-```python
-{code.strip()}
-```
-'''
-
 def generate_vis(df, mode='Document'):
   if len(df) == 0:
     return ''
@@ -50,11 +25,11 @@ def generate_vis(df, mode='Document'):
     if mode == 'Query':
       tok_scores = {m.group(2): float(m.group(1)) for m in re.finditer(r'combine:0=([0-9.]+)\(([^)]+)\)', row.query)}
       max_score = max(tok_scores.values())
-      orig_tokens = factory.tokenizer.tokenize(row.query_0)
+      orig_tokens = factory_max.tokenizer.tokenize(row.query_0)
       id = row.qid
     else:
       tok_scores = row.toks
-      orig_tokens = factory.tokenizer.tokenize(row.text)
+      orig_tokens = factory_max.tokenizer.tokenize(row.text)
       id = row.docno
     def toks2span(toks):
       return '<kbd> </kbd>'.join(f'<kbd style="background-color: rgba(66, 135, 245, {tok_scores.get(t, 0)/max_score});">{t}</kbd>' for t in toks)
@@ -71,170 +46,68 @@ def generate_vis(df, mode='Document'):
 ''')
   return '\n'.join(result)
 
-def predict_query(input):
+def predict_query(input, agg):
   code = f'''import pandas as pd
 import pyterrier as pt ; pt.init()
 import pyt_splade
 
-factory = pyt_splade.SpladeFactory()
+factory = pyt_splade.SpladeFactory(agg={agg})
 
 query_pipeline = factory.query()
 
 query_pipeline({df2code(input)})
 '''
-  res = pipe_queries(input)
+  pipeline = {
+    'max': factory_max,
+    'sum': factory_sum
+  }[agg].query()
+  res = pipeline(input)
   vis = generate_vis(res, mode='Query')
-  return (res, code2md(code), vis)
+  return (res, code2md(code, COLAB_INSTALL, COLAB_NAME), vis)
 
-def predict_doc(input):
+def predict_doc(input, agg):
   code = f'''import pandas as pd
 import pyterrier as pt ; pt.init()
 import pyt_splade
 
-factory = pyt_splade.SpladeFactory()
+factory = pyt_splade.SpladeFactory(agg={agg})
 
 doc_pipeline = factory.indexing()
 
 doc_pipeline({df2code(input)})
 '''
-  res = pipe_docs(input)
+  pipeline = {
+    'max': factory_max,
+    'sum': factory_sum
+  }[agg].indexing()
+  res = pipeline(input)
   vis = generate_vis(res, mode='Document')
   res['toks'] = [json.dumps({k: round(v, 4) for k, v in t.items()}) for t in res['toks']]
-  return (res, code2md(code), vis)
-
-with gr.Blocks(css="table.font-mono td, table.font-mono th { white-space: pre-line; font-size: 11px; line-height: 16px; } table.font-mono td input { width: 95%; } th .cursor-pointer {display: none;} th .min-h-\[2\.3rem\] {min-height: inherit;}") as demo:
-    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>🐕 PyTerrier: SPLADE</h1>")
-    gr.Markdown(open('README.md', 'rt').read().split('\n---\n')[-1])
-
-    example_inp = pd.DataFrame([
+  return (res, code2md(code, COLAB_INSTALL, COLAB_NAME), vis)
+
+interface(
+  MarkdownFile('README.md'),
+  MarkdownFile('query.md'),
+  Demo(
+    predict_query,
+    pd.DataFrame([
       {'qid': '1112389', 'query': 'what is the county for grand rapids, mn'},
-    ])
-    example_out = predict_query(example_inp)
-    inputs, outputs = [], []
-    with gr.Row().style(equal_height=False):
-      with gr.Column(scale=1):
-        with gr.Tab('Pipeline Input'):
-          inputs.append(gr.Dataframe(
-            headers=["qid", "query"],
-            datatype=["str", "str"],
-            col_count=(2, "fixed"),
-            row_count=1,
-            wrap=True,
-            value=example_inp,
-          ))
-        submit_btn = gr.Button("Submit", variant="primary")
-      with gr.Column(scale=2):
-        with gr.Tab('Pipeline Output'):
-          outputs.append(gr.Dataframe(
-            headers=["qid", "query", "docno", "score", "rank", "text"],
-            datatype=["str", "str", "str", "number", "number", "str"],
-            col_count=6,
-            row_count=1,
-            wrap=True,
-            value=example_out[0],
-          ))
-        with gr.Tab('Code'):
-          outputs.append(gr.Markdown(value=example_out[1]))
-        with gr.Tab('Visualisation'):
-          outputs.append(gr.HTML(value=example_out[2]))
-    submit_btn.click(predict_query, inputs, outputs, api_name="predict_query", scroll_to_output=True)
-
-    gr.Markdown('''
-### Document Encoding
-
-The document encoder works similarly to the query encoder: it is a `D→D` (document rewriting, doc-to-doc) transformer, and can be used in pipelines accordingly.
-It maps a document's text into a dictionary with terms from the document re-weighted and weighted expansion terms added.
-
-<div class="pipeline">
-  <div class="df" title="Document Frame">D</div>
-  <div class="transformer" title="SPLADE Indexing Transformer">SPLADE</div>
-  <div class="df" title="Document Frame">D</div>
-</div>
-
-''')
-
-    example_inp = pd.DataFrame([
+    ]),
+    [
+      gr.Dropdown(choices=['max', 'sum'], value='max', label='Aggregation'),
+    ],
+    scale=2/3
+  ),
+  MarkdownFile('doc.md'),
+  Demo(
+    predict_doc,
+    pd.DataFrame([
       {'docno': '0', 'text': 'The presence of communication amid scientific minds was equally important to the success of the Manhattan Project as scientific intellect was. The only cloud hanging over the impressive achievement of the atomic researchers and engineers is what their success truly meant; hundreds of thousands of innocent lives obliterated.'},
-    ])
-    example_out = predict_doc(example_inp)
-    inputs, outputs = [], []
-    with gr.Row().style(equal_height=False):
-      with gr.Column(scale=1):
-        with gr.Tab("Pipeline Input"):
-          inputs.append(gr.Dataframe(
-            headers=["docno", "text"],
-            datatype=["str", "str"],
-            col_count=(2, "fixed"),
-            row_count=1,
-            wrap=True,
-            value=example_inp,
-          ))
-        submit_btn = gr.Button("Submit", variant="primary")
-      with gr.Column(scale=2):
-        with gr.Tab("Pipeline Output"):
-          outputs.append(gr.Dataframe(
-            headers=["qid", "query", "docno", "score", "rank", "text"],
-            datatype=["str", "str", "str", "number", "number", "str"],
-            col_count=6,
-            row_count=1,
-            wrap=True,
-            value=example_out[0],
-          ))
-        with gr.Tab('Code'):
-          outputs.append(gr.Markdown(value=example_out[1]))
-        with gr.Tab('Visualisation'):
-          outputs.append(gr.HTML(value=example_out[2]))
-    submit_btn.click(predict_doc, inputs, outputs, api_name="predict_doc", scroll_to_output=True)
-
-    gr.Markdown('''
-### Putting it all together
-
-When you use the document encoder in an indexing pipeline, the rewritting document contents are indexed:
-
-<div class="pipeline">
-  <div class="df" title="Document Frame">D</div>
-  <div class="transformer" title="SPLADE Indexing Transformer">SPLADE</div>
-  <div class="df" title="Document Frame">D</div>
-  <div class="transformer boring" title="Indexer">Indexer</div>
-  <div class="artefact" title="SPLADE Index">IDX</div>
-</div>
-
-```python
-import pyterrer as pt
-pt.init(version='snapshot')
-import pyt_splade
-
-dataset = pt.get_dataset('irds:msmarco-passage')
-factory = pyt_splade.SpladeFactory()
-
-indexer = pt.IterDictIndexer('./msmarco_psg', pretokenized=True)
-
-indxer_pipe = factory.indexing() >> indexer
-indxer_pipe.index(dataset.get_corpus_iter())
-```
-
-Once you built an index, you can build a retrieval pipeline that first encodes the query,
-and then performs retrieval:
-
-<div class="pipeline">
-  <div class="df" title="Query Frame">Q</div>
-  <div class="transformer" title="SPLADE Query Transformer">SPLADE</div>
-  <div class="df" title="Query Frame">Q</div>
-  <div class="transformer boring" title="Term Frequency Transformer">TF Retriever <div class="artefact" title="SPLADE Index">IDX</div></div>
-  <div class="df" title="Result Frame">R</div>
-</div>
-
-```python
-splade_retr = factory.query() >> pt.BatchRetrieve('./msmarco_psg', wmodel='Tf')
-```
-
-### References & Credits
-
-This package uses [Naver's SPLADE repository](https://github.com/naver/splade).
-
- - Thibault Formal, Benjamin Piwowarski, Stéphane Clinchant. [SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking](https://arxiv.org/abs/2107.05720). SIGIR 2021.
- - Craig Macdonald, Nicola Tonellotto, Sean MacAvaney, Iadh Ounis. [PyTerrier: Declarative Experimentation in Python from BM25 to Dense Retrieval](https://dl.acm.org/doi/abs/10.1145/3459637.3482013). CIKM 2021.
-''')
-
-
-demo.launch(share=False)
+    ]),
+    [
+      gr.Dropdown(choices=['max', 'sum'], value='max', label='Aggregation'),
+    ],
+    scale=2/3
+  ),
+  MarkdownFile('wrapup.md'),
+).launch(share=True)

doc.md

@@ -0,0 +1,10 @@
+### Document Encoding
+
+The document encoder works similarly to the query encoder: it is a `D→D` (document rewriting, doc-to-doc) transformer, and can be used in pipelines accordingly.
+It maps a document's text into a dictionary with terms from the document re-weighted and weighted expansion terms added.
+
+<div class="pipeline">
+  <div class="df" title="Document Frame">D</div>
+  <div class="transformer" title="SPLADE Indexing Transformer">SPLADE</div>
+  <div class="df" title="Document Frame">D</div>
+</div>

query.md

@@ -0,0 +1,11 @@
+### Query Encoding
+
+Let's start by exploring SPLADE's query encoder. The query encoder is a `Q→Q` (query rewriting, query-to-query) transformer, and can be used in pipelines accordingly.
+It maps a query string into [MatchOp](https://terrier-core.readthedocs.io/en/latest/querylanguage.html#matching-op-query-language) query with terms from the
+query re-weighted and weighted expansion terms added.
+
+<div class="pipeline">
+  <div class="df" title="Query Frame">Q</div>
+  <div class="transformer" title="SPLADE Query Transformer">SPLADE</div>
+  <div class="df" title="Query Frame">Q</div>
+</div>

requirements.txt

@@ -1,4 +1,5 @@
 torch
 python-terrier
+git+https://github.com/seanmacavaney/pyterrier_gradio
 git+https://github.com/naver/splade
 git+https://github.com/seanmacavaney/pyt_splade@misc

wrapup.md

@@ -0,0 +1,47 @@
+### Putting it all together
+
+When you use the document encoder in an indexing pipeline, the rewritting document contents are indexed:
+
+<div class="pipeline">
+  <div class="df" title="Document Frame">D</div>
+  <div class="transformer" title="SPLADE Indexing Transformer">SPLADE</div>
+  <div class="df" title="Document Frame">D</div>
+  <div class="transformer boring" title="Indexer">Indexer</div>
+  <div class="artefact" title="SPLADE Index">IDX</div>
+</div>
+
+```python
+import pyterrer as pt
+pt.init(version='snapshot')
+import pyt_splade
+
+dataset = pt.get_dataset('irds:msmarco-passage')
+factory = pyt_splade.SpladeFactory()
+
+indexer = pt.IterDictIndexer('./msmarco_psg', pretokenized=True)
+
+indxer_pipe = factory.indexing() >> indexer
+indxer_pipe.index(dataset.get_corpus_iter())
+```
+
+Once you built an index, you can build a retrieval pipeline that first encodes the query,
+and then performs retrieval:
+
+<div class="pipeline">
+  <div class="df" title="Query Frame">Q</div>
+  <div class="transformer" title="SPLADE Query Transformer">SPLADE</div>
+  <div class="df" title="Query Frame">Q</div>
+  <div class="transformer boring" title="Term Frequency Transformer">TF Retriever <div class="artefact" title="SPLADE Index">IDX</div></div>
+  <div class="df" title="Result Frame">R</div>
+</div>
+
+```python
+splade_retr = factory.query() >> pt.BatchRetrieve('./msmarco_psg', wmodel='Tf')
+```
+
+### References & Credits
+
+This package uses [Naver's SPLADE repository](https://github.com/naver/splade).
+
+ - Thibault Formal, Benjamin Piwowarski, Stéphane Clinchant. [SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking](https://arxiv.org/abs/2107.05720). SIGIR 2021.
+ - Craig Macdonald, Nicola Tonellotto, Sean MacAvaney, Iadh Ounis. [PyTerrier: Declarative Experimentation in Python from BM25 to Dense Retrieval](https://dl.acm.org/doi/abs/10.1145/3459637.3482013). CIKM 2021.