SPLADE

Implementing SPLADE for Efficient Information Retrieval: A PyTorch Guide

In the rapidly evolving field of information retrieval, the quest for models that balance efficiency with effectiveness is ongoing. One of the standout models that has garnered attention for its novel approach is SPLADE - Sparse Lexical and Expansion Model. This model intriguingly combines the depth of neural networks with the efficiency of sparse representations, making it a topic worth exploring. Today, I'll walk you through a simplified example of implementing SPLADE using PyTorch, aiming for clarity and content richness over jargon.

The Essence of SPLADE

At its core, SPLADE leverages the power of transformer-based models, like BERT, to generate sparse vectors for text. These vectors are not only efficient for computation but also effective in capturing the nuanced semantics of language, a crucial aspect for information retrieval tasks. The magic of SPLADE lies in its training process, where it employs a sparsity-inducing loss function. This encourages the model to zero in on the most relevant tokens, leaving behind a trail of zeroes for the rest.

Building SPLADE in PyTorch

Let's dive into the practical side. Our journey begins with the implementation of SPLADE atop PyTorch, a popular deep learning library known for its flexibility and ease of use. Here's a distilled guide to bringing SPLADE to life in your projects.

The SPLADE Model

Our model starts with a familiar face - BERT, leveraging its pre-trained might to understand language. We then introduce a linear layer tasked with scoring the relevance of text, focusing on the essence of what makes or breaks a match in information retrieval.

import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer

class SPLADE(nn.Module):
    def __init__(self, pretrained_model_name='bert-base-uncased'):
        super(SPLADE, self).__init__()
        self.bert = BertModel.from_pretrained(pretrained_model_name)
        self.fc = nn.Linear(768, 1)  # A direct path to scoring relevance

Training with a Twist

Training SPLADE is where things get interesting. Our goal is not just to teach the model about relevance but also to embrace sparsity. This dual focus requires a careful balance, achieved through a specialized loss function.

def train_splade(model, data_loader, optimizer, device):
    model.train()
    for batch in data_loader:
        # Standard training loop, with a SPLADE-specific twist in the loss calculation
        scores = model(input_ids, attention_mask).squeeze()
        loss = compute_loss(scores, labels)  # A concoction of relevance and sparsity

def compute_loss(scores, labels):
    relevance_loss = nn.BCEWithLogitsLoss()(scores, labels)
    sparsity_loss = torch.norm(scores, p=1)  # L1 norm for sparsity
    return relevance_loss + 0.01 * sparsity_loss  # The delicate dance of balancing

The Takeaway

Implementing SPLADE from scratch in PyTorch illuminates the intriguing blend of neural depth and sparse efficiency. This example, while simplified, captures the essence of SPLADE's approach to information retrieval. It's a testament to the model's innovative leveraging of sparsity-inducing techniques, ensuring that every token counts in the vast sea of information.

Diving into SPLADE opens up new vistas for tackling the challenges of information retrieval, providing a framework that's both efficient and effective. As you embark on integrating SPLADE into your projects, remember that the true power of this model lies in its nuanced balance, a reminder of the intricate dance between precision and practicality in the digital age. 

Cross Encoders

A cross-encoder in the context of information retrieval is a type of model architecture designed to enhance the performance of retrieving the most relevant information from a large dataset based on a given query. Unlike traditional or simpler encoder models that process queries and documents separately, cross-encoders evaluate the relevance of a document to a query by jointly encoding both the query and the document together in a single pass. This approach allows the model to consider the intricate interactions between the query and the document, leading to a more nuanced understanding and often superior retrieval performance.

Key characteristics:

• Joint Encoding: Cross-encoders take both the query and a candidate document as input and combine them into a single input sequence, often with a special separator token in between. This allows the model to directly learn interactions between the query and the document.
• Fine-Grained Understanding: By considering the query and document together, cross-encoders can better capture the nuances of relevance, including context, semantic similarities, and specific details that might be missed when encoding them separately.
• Computational Intensity: While providing high accuracy, cross-encoders are computationally more intensive than other architectures like bi-encoders, because each query-document pair must be processed together. This can make them less efficient for applications that require scanning through very large datasets.
• Use in Ranking: Cross-encoders are particularly useful for the ranking stage of information retrieval, where a smaller subset of potentially relevant documents (pre-filtered by a more efficient method, like a bi-encoder) needs to be ranked accurately according to their relevance to the query.
• Training and Fine-tuning: Cross-encoders can be trained or fine-tuned on specific tasks or datasets, allowing them to adapt to the nuances of different domains or types of information retrieval tasks.

In practice, cross-encoders are often used in combination with other models in a two-step retrieval and ranking process. An initial, more efficient model (such as a bi-encoder) quickly narrows down the search space to a manageable number of candidate documents, and then a cross-encoder is applied to this subset to accurately rank the documents in order of relevance to the query. This approach balances the need for both efficiency and high accuracy in information retrieval systems.

Pseudo code:

import torch
from transformers import BertTokenizer, BertForSequenceClassification

class CrossEncoder(torch.nn.Module):
    def __init__(self, pretrained_model_name='bert-base-uncased'):
        super(CrossEncoder, self).__init__()
        self.tokenizer = BertTokenizer.from_pretrained(pretrained_model_name)
        # Assuming a binary classification model where 1 indicates relevance.
        self.model = BertForSequenceClassification.from_pretrained(pretrained_model_name, num_labels=2)

    def forward(self, query, document):
        # Tokenize query and document together, separating them with a [SEP] token
        inputs = self.tokenizer.encode_plus(query, document, return_tensors='pt', add_special_tokens=True, truncation=True, max_length=512)
        # Forward pass through the model
        outputs = self.model(**inputs)
        # Get the logits
        logits = outputs.logits
        # Apply softmax to get probabilities
        probabilities = torch.softmax(logits, dim=1)
        # Assuming label 1 corresponds to "relevant"
        relevance_probability = probabilities[:, 1]
        return relevance_probability

# Example usage
cross_encoder = CrossEncoder()

# Example query and document
query = "What is artificial intelligence?"
document = "Artificial intelligence is a branch of computer science that aims to create intelligent machines."

# Compute relevance score
relevance_score = cross_encoder(query, document)

print(f"Relevance Score: {relevance_score.item()}")

Download a Wikipedia Category

import wikipediaapi
import logging 
import json
from pathlib import Path
from slugify import slugify
import click

logger = logging.getLogger(__name__)

wiki_wiki = wikipediaapi.Wikipedia('parasort (tom-010@web.de)', 'en')

logging.basicConfig()
logging.getLogger().setLevel(logging.INFO)

@click.command()
@click.argument('category', type=str)
@click.option('--out', type=click.Path(path_type=Path), default='out', help='Output directory')
def main(category:str, out:Path):
    download_wikipedia_category(category, out)


def save_page(page, path:Path, breadcumbs: list[str]):
    """Save the content of a Wikipedia page into a file."""

    logger.info(f'Saving page: {page.title}')

    path.mkdir(parents=True, exist_ok=True)
    target = path / f'{slugify(page.title)}.json'
    if target.exists():
        logger.info(f'Page already exists: {page.title}')
        return

    page_json = {
        'title': page.title,
        'summary': page.summary,
        'pageid': page.pageid,
        'breadcumbs': breadcumbs,
        'sections': []
    }
    for section in page.sections:
        page_json['sections'].append({
            'level': section.level,
            'title': section.title,
            'text': section.full_text()
        })
    
    content = page.text
    if content:
        with target.open('w', encoding='utf-8') as f:
            json.dump(page_json, f, ensure_ascii=False, indent=2)

def process_category(category, path:Path, breadcumbs: list[str] = None):
    """Process each category and its subcategories."""
    if breadcumbs is None:
        breadcumbs = []

    logger.info(f'Processing category: {category}')

    for c in category.categorymembers.values():
        if c.ns == wikipediaapi.Namespace.CATEGORY:
            # Create a directory for the subcategory
            category_name = c.title.replace('Category:', '')
            breadcumbs.append(category_name)
            sub_path = path / slugify(category_name)
            process_category(c, sub_path, breadcumbs)
        else:
            save_page(c, path, breadcumbs)

def download_wikipedia_category(category_name:str, out_dir:Path):
    """Download all pages from a Wikipedia category and its subcategories."""
    cat = wiki_wiki.page("Category:" + category_name)
    target = out_dir / slugify(category_name)
    process_category(cat, target)

if __name__ == '__main__':
    main()

Strategy for my keyboard layout

1. Write down all keys on my existing keyboard
2. Build a realistic corpora from a lot of documents (manly own but also similar, + keylogger)
3. Specify how many keys I have
4. Count the occurrences of each key and sort them
5. I have more keys than single characters. Search for patterns, like in sentence piece.

Python: Batch an Interator

from itertools import islice

def batched(iterator, batch_size):
    """
    Takes an iterator and a batch-size, and returns an iterator that yields batches of the given size.
    
    Parameters:
    - iterator: An iterator from which to generate batches.
    - batch_size: An integer specifying the size of each batch.
    
    Yields:
    - Batches of elements from the input iterator, each batch being a list of elements up to the specified batch_size.
    """
    iterator = iter(iterator)  # Ensure it's an iterator
    while True:
        batch = list(islice(iterator, batch_size))
        if not batch:
            break
        yield batch

Usage:

for batch in batched(range(80), 9):
    print(batch)


d = {str(idx): idx for idx in range(10)}

for batch in batched(d.items(), 3):
    for key, value in batch:
        print(key, value)

Output

[0, 1, 2, 3, 4, 5, 6, 7, 8]
[9, 10, 11, 12, 13, 14, 15, 16, 17]
[18, 19, 20, 21, 22, 23, 24, 25, 26]
[27, 28, 29, 30, 31, 32, 33, 34, 35]
[36, 37, 38, 39, 40, 41, 42, 43, 44]
[45, 46, 47, 48, 49, 50, 51, 52, 53]
[54, 55, 56, 57, 58, 59, 60, 61, 62]
[63, 64, 65, 66, 67, 68, 69, 70, 71]
[72, 73, 74, 75, 76, 77, 78, 79]
0 0
1 1
2 2
3 3
4 4
5 5
6 6
7 7
8 8
9 9

Skill estimation and recommender systems

In a recommender system, you try to predict empty cells in the matrix between users and items.

In skill estimation you try to predict empty cell in the matrix between players and games.

Collaborative filtering also has a nice intuition.

pdict: Persistent dict (python)

from sqlitedict import SqliteDict
import sqlite3

def pdict(db_path:str|Path):
    if not isinstance(db_path, Path):
        db_path = Path(db_path)
    if not db_path.exists():
        db_path.parent.mkdir(exist_ok=True, parents=True)
        conn = sqlite3.connect(db_path)
        conn.execute('PRAGMA journal_mode=WAL;')
        conn.commit()
        conn.close()
    return SqliteDict(db_path, autocommit=True)