from transformers import pipeline

model_id = "vemz/pythia-410m-sft-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 SFTTrainer.

Config

python

sft_config = SFTConfig(
    output_dir=output_dir,
    dataset_text_field="text",
    max_length=512,
    packing=True,            
    learning_rate=1e-4,
    per_device_train_batch_size=4, 
    gradient_accumulation_steps=4, 
    num_train_epochs=1,
    fp16=False,
    bf16=False,
    gradient_checkpointing=True,
    logging_steps=10,
    save_strategy="epoch",
)

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 EleutherAI/pythia-410m-deduped 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}}
}