Views
No views yet
1import json
2from random import randrange
3import torch
4
5from peft import LoraConfig, get_peft_model
6from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
7from peft import PeftModel
8
9model1 = AutoModelForCausalLM.from_pretrained(
10 "sysong11/dapt-kogpt", torch_dtype="auto", device_map="auto"
11)
12
13
14lora_path = "sysong11/dapt-kogpt-sum-adapter"
15model2 = PeftModel.from_pretrained(model1, lora_path, device_map="auto")
16tokenizer = AutoTokenizer.from_pretrained(lora_path)
17
18
19test_data = []
20with open("./datasets/test.json", "rb") as f:
21 for line in f:
22 test_data.append(json.loads(line))
23
24
25prompt_template = """\
26<|im_start|>system
27{system_prompt}<|im_end|>
28<|im_start|>user
29{prompt}<|im_end|>
30<|im_start|>assistant"""
31
32msg = "Q:다음 문서를 요약 하세요, Context:{context}"
33
34ix = randrange(len(test_data))
35print(ix)
36datapoint = test_data[ix]
37ref = test_data[ix]["summary_text"]
38system_prompt = "You are an AI assistant. User will you give you a task. Your goal is to complete the task as faithfully as you can."
39tokens = tokenizer.encode(
40 prompt_template.format(
41 system_prompt=system_prompt,
42 prompt=msg.format(context=datapoint["original_text"]),
43 ),
44 return_tensors="pt",
45).to(device="cuda", non_blocking=True)
46
47gen_tokens = model2.generate(
48 input_ids=tokens,
49 do_sample=False,
50 temperature=0.5,
51 max_length=1024,
52 pad_token_id=63999,
53 eos_token_id=63999,
54)
55inputs = tokenizer.batch_decode([gen_tokens[0][: tokens[0].shape[0]]])[0]
56generated = tokenizer.batch_decode([gen_tokens[0][tokens[0].shape[0] :]])[0].replace(
57 "<|im_end|>", ""
58)
59print(inputs)
60print("generated:")
61print(generated)
62
63