kgreenewald commited on
Commit
c164a23
·
verified ·
1 Parent(s): 7562886

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +18 -12
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
- # remove automatic system prompt
114
- string_to_remove = tokenizer.apply_chat_template(question_chat[0:1], tokenize=False,add_generation_prompt=False)
115
- input_text = input_text[len(string_to_remove):]
116
  #tokenize
117
  inputs = tokenizer(input_text, return_tensors="pt")
118
- output = model_UQ.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=80)
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 = question_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])