app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

# --- Model Loading ---
# Load SPLADE v3 model and tokenizer
# "naver/splade-v3" is a common and robust SPLADE v3 model.
# Make sure you've accepted any user access agreements on its Hugging Face Hub page.
try:
    tokenizer = AutoTokenizer.from_pretrained("naver/splade-cocondenser-selfdistil")
    model = AutoModelForMaskedLM.from_pretrained("naver/splade-cocondenser-selfdistil")
    model.eval() # Set the model to evaluation mode for inference
    print("SPLADE v3 model and tokenizer loaded successfully!")
except Exception as e:
    print(f"Error loading SPLADE model or tokenizer: {e}")
    print("Please ensure you have accepted any user access agreements on the Hugging Face Hub page for 'naver/splade-v3' (https://huggingface.co/naver/splade-v3).")
    print("If the problem persists, check your internet connection or try a different SPLADE model if available.")
    tokenizer = None
    model = None

# --- Core SPLADE Representation Function ---
def get_splade_representation(text):
    if tokenizer is None or model is None:
        return "SPLADE model is not loaded. Please check the console for loading errors."

    # Tokenize the input text
    # return_tensors="pt" ensures PyTorch tensors are returned.
    # padding=True pads to the longest sequence in the batch (though for single input, it's just the input length).
    # truncation=True truncates if the text is too long for the model's max input size.
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)

    # Move inputs to the same device as the model (e.g., CPU or GPU)
    # This is important if you were running on a GPU in a production environment.
    # For Hugging Face Spaces free tier, it's usually CPU.
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    # Get the model's output without calculating gradients (inference mode)
    with torch.no_grad():
        output = model(**inputs)

    # Extract the logits from the model's output.
    # SPLADE uses the masked language modeling head's logits to derive term importance.
    # Apply the SPLADE aggregation function: log(1 + ReLU(logits))
    # This transforms the raw logits into a sparse vector where higher values indicate more importance.
    # The attention_mask ensures we only consider actual tokens, not padding.
    
    # Check if 'logits' is in the output (standard for AutoModelForMaskedLM)
    if hasattr(output, 'logits'):
        # Apply the SPLADE transformation
        # output.logits is typically [batch_size, sequence_length, vocab_size]
        # We need to take the max over the sequence_length dimension to get a [batch_size, vocab_size] vector.
        # inputs.attention_mask.unsqueeze(-1) expands the mask to match vocab_size for element-wise multiplication.
        splade_vector = torch.max(torch.log(1 + torch.relu(output.logits)) * inputs['attention_mask'].unsqueeze(-1), dim=1)[0].squeeze()
    else:
        # Fallback/error message if the output structure is unexpected
        return "Model output structure not as expected for SPLADE. 'logits' not found."

    # Convert the sparse vector to a human-readable format.
    # We only care about the non-zero (or very small) entries, as they represent activated terms.
    
    # Get the indices (token IDs) of the non-zero elements in the SPLADE vector
    # torch.nonzero returns coordinates of non-zero elements. squeeze() removes dimensions of size 1.
    # .cpu().tolist() moves the tensor to CPU and converts to a Python list.
    indices = torch.nonzero(splade_vector).squeeze().cpu().tolist()

    # If it's a single index (e.g., a very short text), make it a list for consistent processing
    if not isinstance(indices, list):
        indices = [indices]
        
    # Get the corresponding values (weights) for these non-zero indices
    values = splade_vector[indices].cpu().tolist()

    # Create a dictionary mapping token ID to its weight
    token_weights = dict(zip(indices, values))

    # Decode token IDs back to actual words/subwords
    # Filter out common special tokens that are not meaningful for retrieval (e.g., [CLS], [SEP], [PAD])
    # You can add more tokens to this list if they appear frequently and are not helpful.
    meaningful_tokens = {}
    for token_id, weight in token_weights.items():
        decoded_token = tokenizer.decode([token_id])
        # Filter out special tokens or very short/noisy tokens
        if decoded_token not in ["[CLS]", "[SEP]", "[PAD]", "[UNK]"] and len(decoded_token.strip()) > 0:
            meaningful_tokens[decoded_token] = weight

    # Sort the meaningful tokens by their weight in descending order
    sorted_representation = sorted(meaningful_tokens.items(), key=lambda item: item[1], reverse=True)

    # Format the output for display
    formatted_output = "SPLADE Representation (Top 20 Terms):\n"
    if not sorted_representation:
        formatted_output += "No significant terms found for this input.\n"
    else:
        for i, (term, weight) in enumerate(sorted_representation):
            if i >= 20: # Limit to top 20 terms for readability
                break
            formatted_output += f"- **{term}**: {weight:.4f}\n"
    
    formatted_output += "\n--- Raw SPLADE Vector Info ---\n"
    formatted_output += f"Total non-zero terms in vector: {len(indices)}\n"
    formatted_output += f"Sparsity: {1 - (len(indices) / tokenizer.vocab_size):.2%}\n" # Calculate sparsity

    return formatted_output

# --- Gradio Interface Setup ---
demo = gr.Interface(
    fn=get_splade_representation,
    inputs=gr.Textbox(
        lines=5,
        label="Enter your query or document text here:",
        placeholder="Why is Padua the nicest city in Italy?"
    ),
    outputs=gr.Markdown(), # Use Markdown for richer text formatting (bolding terms)
    title="🌌 SPLADE v3 Sparse Representation Generator",
    description="Enter any text (query or document) to see its SPLADE v3 sparse vector representation. The output highlights the most important terms with their learned weights.",
    allow_flagging="never" # Disable flagging for this demo
)

# Launch the Gradio app
demo.launch()