Views
No views yet
1from transformers import AutoTokenizer
2from trlx.models.modeling_ppo import AutoModelForCausalLMWithHydraValueHead
3
4model = AutoModelForCausalLM.from_pretrained("reciprocate/ppo_hh_gpt-j", revision="checkpoint_05000")
5## there are 01000-02000-03000-...-10000 checkpoints total
6## if you need access to PPO value head, use:
7# model = AutoModelForCausalLMWithHydraValueHead.from_pretrained("reciprocate/ppo_hh_gpt-j", revision="checkpoint_05000")
8## the base model for reference:
9# model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
10
11tokenizer = AutoTokenizer.from_pretrained("gpt2")
12tokenizer.pad_token = tokenizer.eos_token
13tokenizer.padding_side = "left"
14
15prompt_1 = """\
16Human: Hello, can you help me?
17Assistant: Sure, what can I do for you?
18Human: I'm looking for a good recipe for a strawberry cake. What ingredients do I need?
19Assistant:\
20"""
21prompt_2 = """\
22Human: Hi! What kind of music do you like?
23Assistant: I like all kinds of music.
24Human: I'm trying to learn how to play the guitar. Do you have any tips?
25Assistant:\
26"""
27prompts = [prompt_1, prompt_2]
28inputs = tokenizer(
29 [prompt_1, prompt_2],
30 return_tensors="pt",
31 padding=True,
32)
33
34samples = model.generate(
35 **inputs,
36 max_new_tokens=64,
37 top_k=0,
38 top_p=1.0,
39 do_sample=True,
40)
41
42responses = []
43prompt_tokens_lengths = [len(tokenizer.encode(prompt)) for prompt in [prompt_1, prompt_2]]
44
45stop_sequences = ["Human:", "human:", "Assistant:", "assistant:"]
46for i, sample in enumerate(samples):
47 response = tokenizer.decode(sample[prompt_tokens_lengths[i]:], skip_special_tokens=True)
48 # Trim off extra dialogue
49 for stop in stop_sequences:
50 stop_i = response.find(stop)
51 if stop_i >= 0:
52 response = response[:stop_i].rstrip()
53 responses.append(response)
54
55print()
56for prompt, response in zip(prompts, responses):
57 print("=" * 40)
58 print(prompt + response)
59 print("=" * 40)
60 print()
61========================================
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.
========================================