llm-adapters

1
from concurrent.futures import ProcessPoolExecutor
2
import queue
3
import subprocess
4
def evaluate(dataset, gpu):
5
    print('*******dataset:', dataset)
6

7
    command = f"CUDA_VISIBLE_DEVICES={gpu} python evaluate.py \
8
               --model LLaMA-7B \
9
               --adapter LoRA \
10
               --dataset {dataset} \
11
               --base_model '../LLM/models/llama-7b-hf' \
12
               --lora_weights './trained_models/llama-lora'"
13

14
    result = subprocess.run(command, shell=True, text=True, capture_output=False)
15
    print(f"Evaluation results for dataset {dataset} on GPU {gpu}:\n{result.stdout}")
16
    return gpu
17

18

19
datasets = ['AQuA', 'AddSub', 'MultiArith', 'SingleEq', 'gsm8k', 'SVAMP']
20
gpus = [1, 2, 3]
21
tasks_queue = queue.Queue()
22
gpu_queue = queue.Queue()
23

24
for gpu in gpus:
25
    gpu_queue.put(gpu)
26
for task in datasets:
27
    tasks_queue.put(task)
28

29
num_processes = min(len(datasets), len(gpus))  # number of processes to run in parallel
30

31
with ProcessPoolExecutor(max_workers=num_processes) as executor:
32
    futures = [executor.submit(evaluate, tasks_queue.get(), gpu_queue.get()) for i in range(num_processes)]
33
    for future in futures:
34
        gpu_id = future.result()
35
        gpu_queue.put(gpu_id)
36
        if tasks_queue.qsize() > 0:
37
            futures.append(executor.submit(evaluate, tasks_queue.get(), gpu_queue.get()))
38

39

40

41

42

43

44