Modelpublic

danielfein/meta-llama_Llama-3.1-8B_bt_reward_template_1205

sourceHugging Facellama3.1updated 4mo agoView on Hugging Face

0likes19downloads

Model Card

Llama 3.1 8B Creative-Writing Reward Model

This model is a fine-tuned LlamaForSequenceClassification reward model based on meta-llama/Llama-3.1-8B. It was trained with TRL reward modeling and should be used to score stories, not to generate text.

Usage

This is a reward model, not a text-generation model. Load it with AutoModelForSequenceClassification and score the story directly as raw text. Do not apply a chat template or wrap the story in a prompt.

python

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_id = "danielfein/meta-llama_Llama-3.1-8B_bt_reward_template_1205"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
if model.config.pad_token_id is None:
    model.config.pad_token_id = tokenizer.pad_token_id

def reward(story: str) -> float:
    inputs = tokenizer(
        story.strip(),
        return_tensors="pt",
        truncation=True,
        max_length=4096,
    ).to(model.device)
    with torch.inference_mode():
        return model(**inputs).logits.squeeze(-1).float().item()

chosen_score = reward(chosen_story)
rejected_score = reward(rejected_story)
print(chosen_score > rejected_score)

Training procedure

Trained with TRL reward modeling.

Weights & Biases run

Framework versions

—TRL: 0.14.0
—Transformers: 4.51.3
—PyTorch: 2.5.1
—Datasets: 3.2.0
—Tokenizers: 0.21.1