GPT-J (with value head weights) trained on HH with PPO following @reciprocated's trlx example here.
- Dataset: Dahoas/full-hh-rlhf
- Logs: https://wandb.ai/sorry/trlx/runs/itvi8qrn
Usage:
from transformers import AutoTokenizer
from trlx.models.modeling_ppo import AutoModelForCausalLMWithHydraValueHead
model = AutoModelForCausalLM.from_pretrained("reciprocate/ppo_hh_gpt-j", revision="checkpoint_05000")
## there are 01000-02000-03000-...-10000 checkpoints total
## if you need access to PPO value head, use:
# model = AutoModelForCausalLMWithHydraValueHead.from_pretrained("reciprocate/ppo_hh_gpt-j", revision="checkpoint_05000")
## the base model for reference:
# model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"
prompt_1 = """\
Human: Hello, can you help me?
Assistant: Sure, what can I do for you?
Human: I'm looking for a good recipe for a strawberry cake. What ingredients do I need?
Assistant:\
"""
prompt_2 = """\
Human: Hi! What kind of music do you like?
Assistant: I like all kinds of music.
Human: I'm trying to learn how to play the guitar. Do you have any tips?
Assistant:\
"""
prompts = [prompt_1, prompt_2]
inputs = tokenizer(
[prompt_1, prompt_2],
return_tensors="pt",
padding=True,
)
samples = model.generate(
**inputs,
max_new_tokens=64,
top_k=0,
top_p=1.0,
do_sample=True,
)
responses = []
prompt_tokens_lengths = [len(tokenizer.encode(prompt)) for prompt in [prompt_1, prompt_2]]
stop_sequences = ["Human:", "human:", "Assistant:", "assistant:"]
for i, sample in enumerate(samples):
response = tokenizer.decode(sample[prompt_tokens_lengths[i]:], skip_special_tokens=True)
# Trim off extra dialogue
for stop in stop_sequences:
stop_i = response.find(stop)
if stop_i >= 0:
response = response[:stop_i].rstrip()
responses.append(response)
print()
for prompt, response in zip(prompts, responses):
print("=" * 40)
print(prompt + response)
print("=" * 40)
print()
Output:
========================================
Human: Hello, can you help me?
Assistant: Sure, what can I do for you?
Human: I'm looking for a good recipe for a strawberry cake. What ingredients do I need?
Assistant: What kind of cake are you looking for? Filled baked? Chewy, creamy?
========================================
========================================
Human: Hi! What kind of music do you like?
Assistant: I like all kinds of music.
Human: I'm trying to learn how to play the guitar. Do you have any tips?
Assistant: I will give you some tips. The first thing is to practice.. Guitar can be frustrating. Being frustrated with yourself. but if you practice. you will get good. You can practice yourself. Only and only practice. Just practice. Everything. Audio. Music. Visual. Everything. practice. practice.
========================================
Model card borrowed from https://huggingface.co/jon-tow/hh-gpt-j