|
|
import os |
|
|
import json |
|
|
from tqdm import tqdm |
|
|
from loguru import logger |
|
|
from dataclasses import dataclass, asdict |
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed |
|
|
from models.Base import BaseModel |
|
|
from dataloaders.ProblemState import ProblemState |
|
|
from memories.Memory import BaseMemory |
|
|
|
|
|
|
|
|
|
|
|
class BaseAgent: |
|
|
|
|
|
def __init__(self, model: BaseModel, dataset): |
|
|
self.model = model |
|
|
self.dataset = dataset |
|
|
self.memories = self.memory_init() |
|
|
|
|
|
def memory_init(self): |
|
|
return [BaseMemory(ps) for ps in self.dataset.problem_states] |
|
|
def run_single_pass(self, mem: BaseMemory, verbose=False, temperature=0): |
|
|
pass |
|
|
|
|
|
def run(self, output_path=None, multi_thread=True, verbose=False, datalen=None, mem_path=None, temperature=0): |
|
|
data_len = datalen if datalen else len(self.dataset) |
|
|
with tqdm(total=data_len) as pbar: |
|
|
if multi_thread: |
|
|
thread_num = 3 |
|
|
|
|
|
with ThreadPoolExecutor(max_workers=thread_num) as executor: |
|
|
futures = {executor.submit(self.run_single_pass, mem, temperature): mem for mem in self.memories[:data_len]} |
|
|
for future in as_completed(futures): |
|
|
pbar.update(1) |
|
|
else: |
|
|
for mem in self.memories[:data_len]: |
|
|
self.run_single_pass(mem, verbose, temperature=temperature) |
|
|
pbar.update(1) |
|
|
|
|
|
|
|
|
if output_path is not None: |
|
|
self.dataset.write_file(output_path) |
|
|
|
|
|
if mem_path is not None: |
|
|
self.write_memories(mem_path) |
|
|
|
|
|
def write_memories(self, file_path): |
|
|
with open(file_path, "w") as f: |
|
|
for mem in self.memories: |
|
|
output = asdict(mem) |
|
|
f.write(json.dumps(output) + "\n") |
|
|
|
|
|
class SequentialBaseAgent(BaseAgent): |
|
|
def run(self, output_path=None, multi_thread=True, verbose=False, datalen=None, iteration_num=0, temperature=0): |
|
|
data_len = datalen if datalen else len(self.dataset) |
|
|
|
|
|
for iter in range(iteration_num): |
|
|
logger.info(f"\n=== Iteration {iter} ===") |
|
|
root, extension = os.path.splitext(output_path) |
|
|
iter_path = f"{root}_{iter}{extension}" |
|
|
with tqdm(total=data_len) as pbar: |
|
|
if multi_thread: |
|
|
thread_num = 3 |
|
|
|
|
|
with ThreadPoolExecutor(max_workers=thread_num) as executor: |
|
|
futures = {executor.submit(self.run_single_pass, mem, temperature): mem for mem in self.memories[:data_len]} |
|
|
for future in as_completed(futures): |
|
|
pbar.update(1) |
|
|
|
|
|
else: |
|
|
for mem in self.memories[:data_len]: |
|
|
self.run_single_pass(mem, verbose, temperature=temperature) |
|
|
pbar.update(1) |
|
|
|
|
|
|
|
|
if output_path is not None: |
|
|
self.dataset.write_file(iter_path) |
