| import os |
| import json |
| import torch |
| from tqdm import tqdm |
| from transformers import AutoTokenizer, AutoModel, AutoConfig |
|
|
| model_name = "chatglm" |
|
|
| model_path = "../../../base-model/chatglm-6b" |
| tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) |
| config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) |
| model = AutoModel.from_pretrained(model_path, config=config, trust_remote_code=True) |
|
|
| model = model.quantize(4) |
| model = model.half().cuda() |
| model = model.eval() |
|
|
| input_fname = f"../QiaoBan/test_data.json" |
| output_fname = f"./{model_name}/result.json" |
|
|
| with open(input_fname, "r", encoding="utf-8") as fr, open(output_fname, "w", encoding="utf-8") as fw: |
| for line in tqdm(fr.readlines(), ncols=80): |
| sample = json.loads(line.strip()) |
| prompt = sample["prompt"] |
| history = sample["history"] |
| response, _ = model.chat(tokenizer, prompt, history=history) |
| sample["prediction"] = response |
| fw.write(json.dumps(sample, ensure_ascii=False)) |
| fw.write("\n") |
|
|
