fidbench/case_study.py at main · wenhaoli-xmu/fidbench · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import json
import torch
import json
import numpy as np
import random
import argparse
from fidbench.misc import get_model_and_tokenizer
import transformers
import os
from tqdm import tqdm
transformers.logging.set_verbosity_error()


def seed_everything(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.cuda.manual_seed_all(seed)


if __name__ == '__main__':
    seed_everything(42)

    parser = argparse.ArgumentParser()
    parser.add_argument("--env_conf", type=str, default=None)
    parser.add_argument("--max-gen", type=int, default=1024)
    args = parser.parse_args()

    with open(args.env_conf, "r") as f:
        env_conf = json.load(f)

    output_path = os.path.join("case_study", args.env_conf.split('/')[-1].replace(".json", ".jsonl"))
    generation_kwargs = env_conf['generation_kwargs']
    run_name = args.env_conf.replace('.json', '.jsonl')
    tokenizer, model = get_model_and_tokenizer(**env_conf['model'])

    if os.path.exists(output_path):
        os.remove(output_path)

    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    prompts = []
    with open("case_study.jsonl", 'r') as f:
        for prompt in f:
            prompts.append(prompt)

    with open(output_path, 'w') as wf:
        for prompt in tqdm(prompts):
            chats = json.loads(prompt)['conversations']

            input_ids = tokenizer.apply_chat_template(chats, return_tensors='pt', add_generation_prompt=True)

            outputs = model.generate(
                input_ids=input_ids.cuda(),
                max_new_tokens=args.max_gen,
                eos_token_id=tokenizer.eos_token_id,
                **generation_kwargs)

            ctx_len = input_ids.shape[-1]
            outputs = outputs[:, ctx_len:]
            outputs = tokenizer.decode(outputs.ravel().tolist(), skip_special_tokens=True)

            if hasattr(model, 'reset_mask'):
                model.reset_mask()

            prompt_text = tokenizer.decode(input_ids[0], skip_special_tokens=False)
            output_json = json.dumps({"text": outputs})
            wf.write(output_json + '\n')