from transformers import AutoTokenizer, AutoModel
import torch

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("lei-ucsd/mamba-retriever-10k_data-dist_chunk")
model = AutoModel.from_pretrained("lei-ucsd/mamba-retriever-10k_data-dist_chunk")

# Example usage for information retrieval
text = "Your query text here"
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)

with torch.no_grad():
    outputs = model(**inputs)
    # Use binary_logits for relevance scoring
    binary_logits = outputs.binary_logits

Model Architecture

Base Model: state-spaces/mamba2-130m
Training Dataset: lei-ucsd/10k_data
Configuration: dist_chunk
Task: Binary classification for passage relevance in information retrieval

The model extends the base Mamba architecture with a binary classification head for passage relevance scoring.

Training Data

This model was trained on the lei-ucsd/10k_data dataset with dist_chunk chunking strategy.

Limitations

This model is designed specifically for information retrieval tasks and may not perform well on other text classification tasks.

Citation

If you use this model, please cite the original Mamba paper and the dataset used for training.