import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_id = "yophis/DRM-T5-Base-storycloze"

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token

# Load the model
model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto")
model.config.pad_token_id = model.config.eos_token_id

# Input template
input_text = "Story continuation choices: (A) {choice1} (B) {choice2} Story: {sentences}"

Training Details

Training Data

We finetune the model on StoryCloze dataset.

Training Hyperparameters

Learning Rate: 1e-4
Weight Decay: 0.0
Training Steps: 50000
Batch Size: 1024
Precision: bf16 mixed precision

Citation

If you find this model useful, please consider citing our paper:

bibtex

@article{chaichana2025decom,
  title={Decom-Renorm-Merge: Model Merging on the Right Space Improves Multitasking},
  author={Chaichana, Yuatyong and Trachu, Thanapat and Limkonchotiwat, Peerat and Preechakul, Konpat and Khandhawit, Tirasan and Chuangsuwanich, Ekapol},
  journal={arXiv preprint arXiv:2505.23117},
  year={2025}
}

Please also cite StoryCloze and the original T5 model.