support RM metrics, add generating Args
Former-commit-id: c461c6190bc124e98dde7f3cf96a59ce40b26fb0
This commit is contained in:
@@ -15,7 +15,7 @@ from transformers import TextIteratorStreamer
|
||||
|
||||
def main():
|
||||
|
||||
model_args, data_args, finetuning_args = prepare_infer_args()
|
||||
model_args, data_args, finetuning_args, generating_args = prepare_infer_args()
|
||||
model_name = "BLOOM" if "bloom" in model_args.model_name_or_path else "LLaMA"
|
||||
model, tokenizer = load_pretrained(model_args, finetuning_args)
|
||||
|
||||
@@ -25,17 +25,10 @@ def main():
|
||||
def predict_and_print(query, history: list):
|
||||
input_ids = tokenizer([prompt_template.get_prompt(query, history)], return_tensors="pt")["input_ids"]
|
||||
input_ids = input_ids.to(model.device)
|
||||
gen_kwargs = {
|
||||
"input_ids": input_ids,
|
||||
"do_sample": True,
|
||||
"top_p": 0.7,
|
||||
"temperature": 0.95,
|
||||
"num_beams": 1,
|
||||
"max_new_tokens": 512,
|
||||
"repetition_penalty": 1.0,
|
||||
"logits_processor": get_logits_processor(),
|
||||
"streamer": streamer
|
||||
}
|
||||
gen_kwargs = generating_args.to_dict()
|
||||
gen_kwargs["input_ids"] = input_ids
|
||||
gen_kwargs["logits_processor"] = get_logits_processor()
|
||||
gen_kwargs["streamer"] = streamer
|
||||
thread = Thread(target=model.generate, kwargs=gen_kwargs)
|
||||
thread.start()
|
||||
response = ""
|
||||
|
||||
Reference in New Issue
Block a user