Update README.md
Browse files
README.md
CHANGED
@@ -94,35 +94,37 @@ tokenizer = AutoTokenizer.from_pretrained(BASE_NAME,padding_side='left',trust_re
|
|
94 |
model_base = AutoModelForCausalLM.from_pretrained(BASE_NAME,device_map="auto")
|
95 |
model_UQ = PeftModel.from_pretrained(model_base, LORA_NAME)
|
96 |
|
97 |
-
|
98 |
-
question = "What is IBM?"
|
99 |
print("Question:" + question)
|
100 |
question_chat = [
|
101 |
-
{
|
102 |
-
"role": "system",
|
103 |
-
"content": ""
|
104 |
-
},
|
105 |
{
|
106 |
"role": "user",
|
107 |
"content": question
|
108 |
},
|
109 |
]
|
110 |
|
111 |
-
# Generate answer
|
112 |
input_text = tokenizer.apply_chat_template(question_chat,tokenize=False,add_generation_prompt=True)
|
113 |
-
|
114 |
-
|
115 |
-
input_text = input_text[len(string_to_remove):]
|
116 |
#tokenize
|
117 |
inputs = tokenizer(input_text, return_tensors="pt")
|
118 |
-
output =
|
119 |
output_text = tokenizer.decode(output[0])
|
120 |
answer = output_text.split("assistant<|end_of_role|>")[1]
|
121 |
print("Answer: " + answer)
|
122 |
|
123 |
# Generate certainty score
|
124 |
uq_generation_prompt = "<|start_of_role|>certainty<|end_of_role|>"
|
125 |
-
uq_chat =
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
126 |
{
|
127 |
"role": "assistant",
|
128 |
"content": answer
|
@@ -131,7 +133,11 @@ uq_chat = question_chat + [
|
|
131 |
|
132 |
uq_text = tokenizer.apply_chat_template(uq_chat,tokenize=False) + uq_generation_prompt
|
133 |
# remove automatic system prompt
|
|
|
|
|
134 |
uq_text = uq_text[len(string_to_remove):]
|
|
|
|
|
135 |
inputs = tokenizer(uq_text, return_tensors="pt")
|
136 |
output = model_UQ.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=1)
|
137 |
output_text = tokenizer.decode(output[0])
|
|
|
94 |
model_base = AutoModelForCausalLM.from_pretrained(BASE_NAME,device_map="auto")
|
95 |
model_UQ = PeftModel.from_pretrained(model_base, LORA_NAME)
|
96 |
|
97 |
+
question = "What is IBM Research?"
|
|
|
98 |
print("Question:" + question)
|
99 |
question_chat = [
|
|
|
|
|
|
|
|
|
100 |
{
|
101 |
"role": "user",
|
102 |
"content": question
|
103 |
},
|
104 |
]
|
105 |
|
106 |
+
# Generate answer with base model
|
107 |
input_text = tokenizer.apply_chat_template(question_chat,tokenize=False,add_generation_prompt=True)
|
108 |
+
|
109 |
+
|
|
|
110 |
#tokenize
|
111 |
inputs = tokenizer(input_text, return_tensors="pt")
|
112 |
+
output = model_base.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=600)
|
113 |
output_text = tokenizer.decode(output[0])
|
114 |
answer = output_text.split("assistant<|end_of_role|>")[1]
|
115 |
print("Answer: " + answer)
|
116 |
|
117 |
# Generate certainty score
|
118 |
uq_generation_prompt = "<|start_of_role|>certainty<|end_of_role|>"
|
119 |
+
uq_chat = [
|
120 |
+
{
|
121 |
+
"role": "system",
|
122 |
+
"content": ""
|
123 |
+
},
|
124 |
+
{
|
125 |
+
"role": "user",
|
126 |
+
"content": question
|
127 |
+
},
|
128 |
{
|
129 |
"role": "assistant",
|
130 |
"content": answer
|
|
|
133 |
|
134 |
uq_text = tokenizer.apply_chat_template(uq_chat,tokenize=False) + uq_generation_prompt
|
135 |
# remove automatic system prompt
|
136 |
+
string_to_remove = tokenizer.apply_chat_template(uq_chat[0:1], tokenize=False,add_generation_prompt=False)
|
137 |
+
input_text = input_text[len(string_to_remove):]
|
138 |
uq_text = uq_text[len(string_to_remove):]
|
139 |
+
|
140 |
+
# tokenize and generate
|
141 |
inputs = tokenizer(uq_text, return_tensors="pt")
|
142 |
output = model_UQ.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=1)
|
143 |
output_text = tokenizer.decode(output[0])
|