Views
No views yet
1from transformers import AutoTokenizer
2from trlx.models.modeling_ppo import AutoModelForCausalLMWithHydraValueHead
3
4model = AutoModelForCausalLMWithHydraValueHead.from_pretrained("jon-tow/hh-gpt-j")
5# original_model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
6
7tokenizer = AutoTokenizer.from_pretrained("gpt2")
8tokenizer.pad_token = tokenizer.eos_token
9tokenizer.padding_side = "left"
10
11prompt_1 = """\
12Human: Hello, can you help me?
13Assistant: Sure, what can I do for you?
14Human: I'm looking for a good recipe for a strawberry cake. What ingredients do I need?
15Assistant:\
16"""
17prompt_2 = """\
18Human: Hi! What kind of music do you like?
19Assistant: I like all kinds of music.
20Human: I'm trying to learn how to play the guitar. Do you have any tips?
21Assistant:\
22"""
23prompts = [prompt_1, prompt_2]
24inputs = tokenizer(
25 [prompt_1, prompt_2],
26 return_tensors="pt",
27 padding=True,
28)
29
30samples = model.generate(
31 **inputs,
32 max_new_tokens=64,
33 top_k=0,
34 top_p=1.0,
35 do_sample=True,
36)
37
38responses = []
39prompt_tokens_lengths = [len(tokenizer.encode(prompt)) for prompt in [prompt_1, prompt_2]]
40stop_sequences = ["Human:", "human:", "Assistant:", "assistant:"]
41for i, sample in enumerate(samples):
42 response = tokenizer.decode(sample[prompt_tokens_lengths[i]:], skip_special_tokens=True)
43 # Trim off extra dialogue
44 for stop in stop_sequences:
45 stop_i = response.find(stop)
46 if stop_i >= 0:
47 response = response[:stop_i].rstrip()
48 responses.append(response)
49
50print()
51for prompt, response in zip(prompts, responses):
52 print("=" * 40)
53 print(prompt + response)
54 print("=" * 40)
55 print()========================================
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: Is strawberry flavour a primary flavour you want in the cake?
========================================
========================================
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: One thing you can try is to form chords and strums. Form chords and strums will help you to practice and learn how to play instruments easily. You can also download free music online. Besure to check out different genres and instruments. You don't have to learn everyone all at once. Learning the basics is
========================================