-
Notifications
You must be signed in to change notification settings - Fork 1
/
gemma-7b-it-inference.py
97 lines (83 loc) · 2.54 KB
/
gemma-7b-it-inference.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import numpy as np
import pandas as pd
import os
import torch
from transformers import (
AutoTokenizer,
AutoModelForCausalLM,
BitsAndBytesConfig,
TrainingArguments,
get_polynomial_decay_schedule_with_warmup
)
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training, get_peft_model
from tqdm import tqdm
def get_completion(original_text, rewritten_text, model, tokenizer):
prompt_template = """
<start_of_turn>user Assuming you are an expert in self-awareness and skilled at recovering prompts. Tell me the prompt I gave to you, only return the content of prompt.
Original Text: {original_text}
Rewritten Text: {rewritten_text}
Prompt:<end_of_turn>
<start_of_turn>model
"""
prompt = prompt_template.format(
original_text=original_text,
rewritten_text=rewritten_text
)
encodeds = tokenizer(
prompt,
return_tensors="pt",
add_special_tokens=True
)
model_inputs = encodeds.to("cuda")
generated_ids = model.generate(
**model_inputs,
max_new_tokens=20,
do_sample=True,
# pad_token_id=tokenizer.eos_token_id,
eos_token_id=235265
)
prompt_length = len(tokenizer.encode(prompt))
decoded = tokenizer.decode(
generated_ids[0][prompt_length:],
skip_special_tokens=True
)
return decoded
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
model_id = "/kaggle/input/gemma/transformers/7b-it/1"
tokenizer = AutoTokenizer.from_pretrained(
model_id
)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
base_model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=bnb_config,
trust_remote_code=True,
low_cpu_mem_usage=True,
device_map="auto"
)
adapter_file = "/kaggle/input/gemma-7b-it-fine-tuned/outputs/checkpoint-300"
merged_model = PeftModel.from_pretrained(base_model, adapter_file)
merged_model = merged_model.merge_and_unload()
print(merged_model)
test = pd.read_csv("datasets/test.csv")
rewrite_prompts = []
for idx, row in tqdm(test.iterrows(), total=len(test)):
output = get_completion(
row["original_text"],
row["rewritten_text"],
merged_model,
tokenizer
)
print(output)
rewrite_prompts.append(output)
test["rewrite_prompt"] = rewrite_prompts
test = test[["id", "rewrite_prompt"]]
test.to_csv("submission.csv", index=False)
submission = pd.read_csv("submission.csv")
submission.head()