This model builds on the SFT checkpoint vemz/pythia-410m-sft-imdb, which itself is a LoRA-adapted version of EleutherAI/pythia-410m-deduped. Its objective is to generate positive movie reviews.

It has been trained using TRL.

The final reward model uses lvwerra/distilbert-imdb with penalties.

Quick start

python

from transformers import pipeline

model_id = "vemz/pythia-410m-rloo-imdb"

generator = pipeline("text-generation", model=model_id, device_map="auto")

prompt = "At first I wad bored, but then"

output = generator(prompt, max_new_tokens=50, return_full_text=False)[0]
print(output["generated_text"])

Training procedure

This model was trained with RLOOTrainer.

Config

python

rloo_config = RLOOConfig(
    output_dir=output_dir,
    learning_rate=3e-5,
    lr_scheduler_type="cosine",
    warmup_ratio=0.1,
    num_generations=4,
    beta=0.05,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=16,
    max_steps=100,
    logging_steps=10,
    bf16=False,
    fp16=False,
    gradient_checkpointing=True,
    max_completion_length=48,
    save_strategy="steps",
    save_steps=40,
    save_total_limit=1,
)

Limitations and Risks

This model is not instruction-tuned.
It may generate biased, offensive, or inappropriate language present in the IMDb dataset.
It may hallucinate facts or produce incorrect information.
This model should not be used for safety-critical or factual applications.

Adapter Notice

This repository contains LoRA adapter weights. The base model vemz/pythia-410m-sft-imdb must be loaded separately.

Framework versions

PEFT 0.18.0
TRL: 0.26.2
Transformers: 4.57.3
Pytorch: 2.9.1
Datasets: 4.4.2
Tokenizers: 0.22.1

Citations

Cite TRL as:

bibtex

@misc{vonwerra2022trl,
	title        = {{TRL: Transformer Reinforcement Learning}},
	author       = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec},
	year         = 2020,
	journal      = {GitHub repository},
	publisher    = {GitHub},
	howpublished = {\url{https://github.com/huggingface/trl}}
}