vemz/pythia-410m-sft-imdb

Model Card for sft_imdb_pythia_410m

This model is a fine-tuned version of EleutherAI/pythia-410m-deduped on the IMDb dataset to generate film reviews. It has been trained using TRL.

Quick start

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

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:

@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}}
}