Views
No views yet
1from datasets import load_dataset
2
3from transformers import AutoTokenizer, AutoModelForCausalLM
4
5dataset = load_dataset("CarperAI/openai_summarize_tldr")
6
7val_prompts = [sample["prompt"] for sample in dataset["valid"]]
8
9kwargs = {
10 "max_new_tokens": 50,
11 "do_sample": True,
12 "top_k": 0,
13 "top_p": 1,
14}
15
16model = AutoModelForCausalLM.from_pretrained("pvduy/ppo_pythia6B_sample")
17model.eval()
18tokenizer = AutoTokenizer.from_pretrained("pvduy/ppo_pythia6B_sample")
19tokenizer.pad_token_id = tokenizer.eos_token_id
20
21count = 0
22
23for prompt in val_prompts:
24 output_tk = tokenizer(prompt, return_tensors="pt")
25 outputs = model.generate(output_tk.input_ids, attention_mask=output_tk.attention_mask, **kwargs)
26 print("Prompt:", prompt)
27 print("Output:", tokenizer.decode(outputs[0], skip_special_tokens=True).split("TL;DR:")[1].strip())
28 print("=================================")
29 count += 1
30 if count == 10:
31 break