批量预测问题 about knowlm HOT 13 CLOSED

try:
    if torch.backends.mps.is_available():
        device = "mps"
except:
    pass
def tokenize(prompt, add_eos_token=True):
    """

    :param prompt:
    :param add_eos_token:
    :return:
    """
    return tokenizer(
        prompt,
        return_tensors="pt",
    )

def generate_and_tokenize_prompt(data_point):
    full_prompt = prompter.generate_prompt(
        data_point["instruction"],
        data_point["input"],
    )
    tokenized_full_prompt = tokenize(full_prompt)
    return tokenized_full_prompt

if data_path.endswith(".json") or data_path.endswith(".jsonl"):
    data = load_dataset("json", data_files=data_path)
    print(f"data includes: {data_path}")
else:
    """is folder"""
    data_paths = []
    data_path = data_path if data_path[-1] == "/" else data_path + "/"
    for i in os.listdir(data_path):
        data_paths.append(os.path.join(data_path, i))
    print(f"data includes: {data_paths}")
    data = load_dataset(data_paths)

if val_set_size > 0:
    test_data = data.map(generate_and_tokenize_prompt)
input_ids = test_data.data['train']['input_ids']

if device == "cuda":
    model = LlamaForCausalLM.from_pretrained(
        base_model,
        load_in_8bit=False,
        torch_dtype=torch.float16,
        device_map="auto",
    )
    model = PeftModel.from_pretrained(
        model,
        lora_weights,
        torch_dtype=torch.float16,
    )
    model.config.pad_token_id = tokenizer.pad_token_id = 0  # same as unk token id
    model.config.bos_token_id = tokenizer.bos_token_id = 1
    model.config.eos_token_id = tokenizer.eos_token_id = 2
    print("loaded model")
model.eval()
if torch.__version__ >= "2" and sys.platform != "win32":
    model = torch.compile(model)


generation_config = GenerationConfig(
    temperature=1.0,
    top_p=1.0,
    top_k=50,
    num_beams=1,
    max_new_tokens=256,
    pad_token_id=tokenizer.pad_token_id,
    return_dict_in_generate=True,
    output_scores=True
)
with torch.no_grad():
    outputs = model.generate(input_ids=input_ids, generation_config=generation_config)

# 解码生成的文本
decoded_texts = [tokenizer.decode(output.sequences[0], skip_special_tokens=True) for output in outputs]
with open(output_dir, "w") as writer:
    for i, result in tqdm(enumerate(decoded_texts)):
        print(i, result)
        writer.write(result + '\n')