from datasets import Dataset

from sentence_transformers.sparse_encoder import SparseEncoder, SparseEncoderTrainer, losses

model = SparseEncoder("sparse-embedding/splade-distilbert-base-uncased-init")
guide = SparseEncoder("naver/splade-cocondenser-ensembledistil")
train_dataset = Dataset.from_dict(
    {
        "anchor": ["It's nice weather outside today.", "He drove to work."],
        "positive": ["It's so sunny.", "He took the car to the office."],
    }
)
loss = losses.SparseCachedGISTEmbedLoss(
    model,
    guide,
    mini_batch_size=64,
    margin_strategy="relative",  # or "relative" (e.g., margin=0.05 for max. 95% of positive similarity)
    margin=0.1,
)

trainer = SparseEncoderTrainer(model=model, train_dataset=train_dataset, loss=loss)
trainer.train()

# TODO: Investigate if it's working with a test, seems that the problem is hparam and not the cache part

from datasets import Dataset

from sentence_transformers.sparse_encoder import (
    MLMTransformer,
    SparseCachedGISTEmbedLoss,
    SparseEncoder,
    SparseEncoderTrainer,
    SpladePooling,
)

# Initialize the SPLADE model
model_name = "naver/splade-cocondenser-ensembledistil"
model = SparseEncoder(
    modules=[
        MLMTransformer(model_name),
        SpladePooling(pooling_strategy="max"),  # You can also use 'sum'
    ],
    device="cuda:0",
)

# Create a small toy dataset
train_dataset = Dataset.from_dict(
    {
        "anchor": ["It's nice weather outside today.", "He drove to work."],
        "positive": ["It's so sunny.", "He took the car to the office."],
    }
)

# Initialize the sparse loss with a guide model
guide = SparseEncoder(
    modules=[
        MLMTransformer("prithivida/Splade_PP_en_v1"),
        SpladePooling(pooling_strategy="max"),
    ],
    device="cuda:0",
)
loss = SparseCachedGISTEmbedLoss(model, guide=guide)

# Create the trainer
trainer = SparseEncoderTrainer(model=model, train_dataset=train_dataset, loss=loss)

# Train the model
trainer.train()
