Views
No views yet
LlamaForSequenceClassification reward model for scoring creative-writing stories. It should be used as a scalar verifier/reward model, not as a text-generation model.AutoModelForSequenceClassification and score the story directly as raw text. Do not apply a chat template or wrap the story in a prompt.1import torch
2from transformers import AutoModelForSequenceClassification, AutoTokenizer
3
4model_id = "SAA-Lab/Llama8B-CreativeWritingVerifier"
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6model = AutoModelForSequenceClassification.from_pretrained(
7 model_id,
8 torch_dtype=torch.bfloat16,
9 device_map="auto",
10)
11
12if tokenizer.pad_token is None:
13 tokenizer.pad_token = tokenizer.eos_token
14if model.config.pad_token_id is None:
15 model.config.pad_token_id = tokenizer.pad_token_id
16
17def reward(story: str) -> float:
18 inputs = tokenizer(
19 story.strip(),
20 return_tensors="pt",
21 truncation=True,
22 max_length=4096,
23 ).to(model.device)
24 with torch.inference_mode():
25 return model(**inputs).logits.squeeze(-1).float().item()
26
27chosen_score = reward(chosen_story)
28rejected_score = reward(rejected_story)
29print(chosen_score > rejected_score)