Views
No views yet
Cosine Scheduler
Learning Rate: 9e-6
Warmup Ratio: 0.03
Batch Size: 256
Epoch: 1transformers library to score the quality of a generated response to a given prompt. The input format should match what the model was trained on (e.g., a full conversation turn using the Llama 3 chat template).1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4model_id = "OpenRLHF/Llama-3-8b-rm-mixture" # This model ID
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6# Ensure to load with appropriate torch_dtype, e.g., torch.bfloat16 for Llama models
7model = AutoModelForSequenceClassification.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
8
9# Example: Score responses to a user prompt
10prompt = "Write a short poem about a cat."
11response_good = "A feline friend, soft and sleek,\
12Curled up warm, a purring peek.\
13Through sunlit naps and playful chase,\
14Graceful paws in every space."
15response_bad = "Cats are okay. They sit sometimes. Dog is better."
16
17# Apply the chat template for the full conversation turn (user prompt + assistant response)
18# The `apply_chat_template` method structures the input as expected by the model.
19messages_good = [
20 {"role": "user", "content": prompt},
21 {"role": "assistant", "content": response_good},
22]
23messages_bad = [
24 {"role": "user", "content": prompt},
25 {"role": "assistant", "content": response_bad},
26]
27
28input_ids_good = tokenizer.apply_chat_template(messages_good, return_tensors="pt", add_generation_prompt=False).to(model.device)
29input_ids_bad = tokenizer.apply_chat_template(messages_bad, return_tensors="pt", add_generation_prompt=False).to(model.device)
30
31# Get scores
32with torch.no_grad():
33 score_good = model(input_ids_good).logits.item()
34 score_bad = model(input_ids_bad).logits.item()
35
36print(f"Score for good response: {score_good:.2f}")
37print(f"Score for bad response: {score_bad:.2f}")