initial commit

Dockerfile

@@ -0,0 +1,45 @@
+# Base image
+FROM python:3.10-slim
+
+# Avoid interactive prompts
+ENV DEBIAN_FRONTEND=noninteractive
+
+# System deps (for building Python + Rust crates)
+RUN apt-get update && apt-get install -y \
+    curl \
+    build-essential \
+    default-jre \
+    default-jre-headless \
+    default-jdk \
+    default-jdk-headless \
+    debianutils \
+    git \
+    openssl \
+    libssl-dev \
+    pkg-config \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Rust (latest stable via rustup)
+RUN curl https://sh.rustup.rs -sSf | sh -s -- -y
+ENV PATH="/root/.cargo/bin:${PATH}"
+
+# Verify versions (optional but useful for logs)
+RUN rustc --version && cargo --version
+
+# Set workdir
+WORKDIR /app
+
+# Copy dependency files first (for caching)
+COPY requirements.txt .
+
+# Install Python deps
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the app
+COPY . .
+
+# Expose port (HF expects 7860)
+EXPOSE 7860
+
+# Run the app
+CMD ["python", "app.py"]

README.md

@@ -1,12 +1,16 @@
 ---
-title: Splade3
-emoji: 🌍
-colorFrom: yellow
+title: PyTerrier SPLADE
+emoji: 🐕
+colorFrom: green
 colorTo: green
 sdk: gradio
-sdk_version: 6.10.0
+sdk_version: 5.7.1
+python_version: '3.10'
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🐕 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.

app.py

@@ -0,0 +1,107 @@
+import base64
+import re
+import json
+import pandas as pd
+import gradio as gr
+import pyterrier as pt
+pt.init()
+import pyt_splade
+from pyterrier_gradio import Demo, MarkdownFile, interface, df2code, code2md, EX_Q, EX_D, df2list
+factory_max = pyt_splade.Splade(agg='max')
+factory_sum = pyt_splade.Splade(agg='sum')
+
+COLAB_NAME = 'pyterrier_splade.ipynb'
+COLAB_INSTALL = '''
+!pip install -q git+https://github.com/naver/splade
+!pip install -q git+https://github.com/cmacdonald/pyt_splade
+'''.strip()
+
+def generate_vis(df, mode='Document'):
+  if len(df) == 0:
+    return ''
+  result = []
+  if mode == 'Document':
+    max_score = max(max(t.values()) for t in df['toks'])
+  for row in df.itertuples(index=False):
+    if mode == 'Query':
+      tok_scores = row.query_toks
+      orig_tokens = factory_max.tokenizer.tokenize(row.query)
+      max_score = max(tok_scores.values())
+      id = row.qid
+    else:
+      tok_scores = row.toks
+      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)
+    orig_tokens_set = set(orig_tokens)
+    exp_tokens = [t for t, v in sorted(tok_scores.items(), key=lambda x: (-x[1], x[0])) if t not in orig_tokens_set]
+    result.append(f'''
+<div style="font-size: 1.2em;">{mode}: <strong>{id}</strong></div>
+<div style="margin: 4px 0 16px; padding: 4px; border: 1px solid black;">
+<div>
+{toks2span(orig_tokens)}
+</div>
+<div><strong>Expansion Tokens:</strong> {toks2span(exp_tokens)}</div>
+</div>
+''')
+  return '\n'.join(result)
+
+def predict_query(input, agg):
+  code = f'''import pyt_splade
+
+splade = pyt_splade.Splade(agg={agg!r})
+
+query_pipeline = splade.query_encoder()
+
+query_pipeline({df2list(input)})
+'''
+  pipeline = {
+    'max': factory_max,
+    'sum': factory_sum
+  }[agg].query_encoder()
+  res = pipeline(input)
+  vis = generate_vis(res, mode='Query')
+  res['query_toks'] = [json.dumps({k: round(v, 4) for k, v in t.items()}) for t in res['query_toks']]
+  return (res, code2md(code, COLAB_INSTALL, COLAB_NAME), vis)
+
+def predict_doc(input, agg):
+  code = f'''import pyt_splade
+
+splade = pyt_splade.Splade(agg={repr(agg)})
+
+doc_pipeline = splade.doc_encoder()
+
+doc_pipeline({df2list(input)})
+'''
+  pipeline = {
+    'max': factory_max,
+    'sum': factory_sum
+  }[agg].doc_encoder()
+  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, COLAB_INSTALL, COLAB_NAME), vis)
+
+interface(
+  MarkdownFile('README.md'),
+  MarkdownFile('query.md'),
+  Demo(
+    predict_query,
+    EX_Q,
+    [
+      gr.Dropdown(choices=['max', 'sum'], value='max', label='Aggregation'),
+    ],
+    scale=2/3
+  ),
+  MarkdownFile('doc.md'),
+  Demo(
+    predict_doc,
+    EX_D,
+    [
+      gr.Dropdown(choices=['max', 'sum'], value='max', label='Aggregation'),
+    ],
+    scale=2/3
+  ),
+  MarkdownFile('wrapup.md'),
+).launch(share=False, server_name="0.0.0.0", server_port=7860)

build.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+set -e
+
+# Install rustup (latest stable toolchain)
+curl https://sh.rustup.rs -sSf | sh -s -- -y
+
+# Load cargo into PATH
+source $HOME/.cargo/env
+
+# Ensure latest stable (>= 1.88)
+rustup update stable
+rustup default stable
+
+rustc --version
+cargo --version

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 attn" title="SPLADE Indexing Transformer">SPLADE</div>
+  <div class="df" title="Document Frame">D</div>
+</div>

packages.txt

@@ -0,0 +1,7 @@
+default-jre
+default-jre-headless
+default-jdk
+default-jdk-headless
+debianutils
+build-essential
+curl

query.md

@@ -0,0 +1,10 @@
+### 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 a sparse `dict`-formatted query vector, with the tokens as the keys and the weights as values.
+
+<div class="pipeline">
+  <div class="df" title="Query Frame">Q</div>
+  <div class="transformer attn" title="SPLADE Query Transformer">SPLADE</div>
+  <div class="df" title="Query Frame">Q</div>
+</div>

requirements.txt

@@ -0,0 +1,5 @@
+torch
+python-terrier>=0.11.0
+git+https://github.com/seanmacavaney/pyterrier_gradio@v0.0.9
+git+https://github.com/naver/splade
+git+https://github.com/cmacdonald/pyt_splade

wrapup.md

@@ -0,0 +1,46 @@
+### Putting it all together
+
+When you use the document encoder in an indexing pipeline, the rewritten document contents are indexed:
+
+<div class="pipeline">
+  <div class="df" title="Document Frame">D</div>
+  <div class="transformer attn" title="SPLADE Indexing Transformer">SPLADE</div>
+  <div class="df" title="Document Frame">D</div>
+  <div class="transformer" title="Indexer">Indexer</div>
+  <div class="artefact" title="SPLADE Index">IDX</div>
+</div>
+
+```python
+import pyterrier as pt
+import pyt_splade
+
+dataset = pt.get_dataset('irds:msmarco-passage')
+splade = pyt_splade.Splade()
+
+indexer = pt.IterDictIndexer('./msmarco_psg', pretokenised=True)
+
+indxer_pipe = splade.doc_encoder() >> 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 attn" title="SPLADE Query Transformer">SPLADE</div>
+  <div class="df" title="Query Frame">Q</div>
+  <div class="transformer" 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 = splade.query_encoder() >> pt.terrier.Retriever('./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.