import pandas as pd
from eval.evaluator import Eval
from dataset.base_dataset import DatasetBase
from utils.llm_chain import MetaChain
from estimator import give_estimator
from pathlib import Path
import pickle
import os
import json
import logging
import wandb
class OptimizationPipeline:
"""
The main pipeline for optimization. The pipeline is composed of 4 main components:
1. dataset - The dataset handle the data including the annotation and the prediction
2. annotator - The annotator is responsible generate the GT
3. predictor - The predictor is responsible to generate the prediction
4. eval - The eval is responsible to calculate the score and the large errors
"""
def __init__(self, config, task_description: str = None, initial_prompt: str = None, output_path: str = ''):
"""
Initialize a new instance of the ClassName class.
:param config: The configuration file (EasyDict)
:param task_description: Describe the task that needed to be solved
:param initial_prompt: Provide an initial prompt to solve the task
:param output_path: The output dir to save dump, by default the dumps are not saved
"""
if config.use_wandb: # In case of using W&B
wandb.login()
self.wandb_run = wandb.init(
project="AutoGPT",
config=config,
)
if output_path == '':
self.output_path = None
else:
if not os.path.isdir(output_path):
os.makedirs(output_path)
self.output_path = Path(output_path)
logging.basicConfig(filename=self.output_path / 'info.log', level=logging.DEBUG,
format='%(asctime)s - %(levelname)s - %(message)s', force=True)
self.dataset = None
self.config = config
self.meta_chain = MetaChain(config)
self.initialize_dataset()
self.task_description = task_description
self.cur_prompt = initial_prompt
self.predictor = give_estimator(config.predictor)
self.annotator = give_estimator(config.annotator)
self.eval = Eval(config.eval, self.meta_chain.error_analysis, self.dataset.label_schema)
self.batch_id = 0
self.patient = 0
@staticmethod
def log_and_print(message):
print(message)
logging.info(message)
def initialize_dataset(self):
"""
Initialize the dataset: Either empty dataset or loading an existing dataset
"""
logging.info('Initialize dataset')
self.dataset = DatasetBase(self.config.dataset)
if 'initial_dataset' in self.config.dataset.keys():
logging.info(f'Load initial dataset from {self.config.dataset.initial_dataset}')
self.dataset.load_dataset(self.config.dataset.initial_dataset)
def calc_usage(self):
"""
Calculate the usage of the optimization process (either $ in case of openAI or #tokens the other cases)
"""
total_usage = 0
total_usage += self.meta_chain.calc_usage()
total_usage += self.annotator.calc_usage()
total_usage += self.predictor.calc_usage()
return total_usage
def extract_best_prompt(self):
sorted_history = sorted(
self.eval.history[min(self.config.meta_prompts.warmup - 1, len(self.eval.history) - 1):],
key=lambda x: x['score'],
reverse=False)
return {'prompt': sorted_history[-1]['prompt'], 'score': sorted_history[-1]['score']}
def run_step_prompt(self):
"""
Run the meta-prompts and get new prompt suggestion, estimated prompt score and a set of challenging samples
for the new prompts
"""
step_num = len(self.eval.history)
if (step_num < self.config.meta_prompts.warmup) or (step_num % 3) > 0:
last_history = self.eval.history[-self.config.meta_prompts.history_length:]
else:
sorted_history = sorted(self.eval.history[self.config.meta_prompts.warmup - 1:], key=lambda x: x['score'],
reverse=False)
last_history = sorted_history[-self.config.meta_prompts.history_length:]
history_prompt = '\n'.join([self.eval.sample_to_text(sample,
num_errors_per_label=self.config.meta_prompts.num_err_prompt,
is_score=True) for sample in last_history])
prompt_input = {"history": history_prompt, "task_description": self.task_description,
'error_analysis': last_history[-1]['analysis']}
if 'label_schema' in self.config.dataset.keys():
prompt_input["labels"] = json.dumps(self.config.dataset.label_schema)
prompt_suggestion = self.meta_chain.step_prompt_chain.invoke(prompt_input)
self.log_and_print(f'Previous prompt score:\n{self.eval.mean_score}\n#########\n')
self.log_and_print(f'Get new prompt:\n{prompt_suggestion["prompt"]}')
self.batch_id += 1
if len(self.dataset) < self.config.dataset.max_samples:
batch_input = {"num_samples": self.config.meta_prompts.samples_generation_batch,
"task_description": self.task_description,
"prompt": prompt_suggestion['prompt']}
batch_inputs = self.generate_samples_batch(batch_input, self.config.meta_prompts.num_generated_samples,
self.config.meta_prompts.samples_generation_batch)
if sum([len(t['errors']) for t in last_history]) > 0:
history_samples = '\n'.join([self.eval.sample_to_text(sample,
num_errors_per_label=self.config.meta_prompts.num_err_samples,
is_score=False) for sample in last_history])
for batch in batch_inputs:
extra_samples = self.dataset.sample_records()
extra_samples_text = DatasetBase.samples_to_text(extra_samples)
batch['history'] = history_samples
batch['extra_samples'] = extra_samples_text
else:
for batch in batch_inputs:
extra_samples = self.dataset.sample_records()
extra_samples_text = DatasetBase.samples_to_text(extra_samples)
batch['history'] = 'No previous errors information'
batch['extra_samples'] = extra_samples_text
samples_batches = self.meta_chain.step_samples.batch_invoke(batch_inputs,
self.config.meta_prompts.num_workers)
new_samples = [element for sublist in samples_batches for element in sublist['samples']]
new_samples = self.dataset.remove_duplicates(new_samples)
self.dataset.add(new_samples, self.batch_id)
logging.info('Get new samples')
self.cur_prompt = prompt_suggestion['prompt']
def stop_criteria(self):
"""
Check if the stop criteria holds. The conditions for stopping:
1. Usage is above the threshold
2. There was no improvement in the last > patient steps
"""
if 0 < self.config.stop_criteria.max_usage < self.calc_usage():
return True
if len(self.eval.history) <= self.config.meta_prompts.warmup:
self.patient = 0
return False
min_batch_id, max_score = self.eval.get_max_score(self.config.meta_prompts.warmup-1)
if max_score - self.eval.history[-1]['score'] > -self.config.stop_criteria.min_delta:
self.patient += 1
else:
self.patient = 0
if self.patient > self.config.stop_criteria.patience:
return True
return False
@staticmethod
def generate_samples_batch(batch_input, num_samples, batch_size):
"""
Generate samples in batch
"""
batch_num = num_samples // batch_size
all_batches = [batch_input.copy() for _ in range(batch_num)]
reminder = num_samples - batch_num * batch_size
if reminder > 0:
all_batches.append(batch_input.copy())
all_batches[-1]['num_samples'] = reminder
return all_batches
def generate_initial_samples(self):
"""
In case the initial dataset is empty generate the initial samples
"""
batch_input = {"num_samples": self.config.meta_prompts.samples_generation_batch,
"task_description": self.task_description,
"instruction": self.cur_prompt}
batch_inputs = self.generate_samples_batch(batch_input, self.config.meta_prompts.num_initialize_samples,
self.config.meta_prompts.samples_generation_batch)
samples_batches = self.meta_chain.initial_chain.batch_invoke(batch_inputs, self.config.meta_prompts.num_workers)
samples_list = [element for sublist in samples_batches for element in sublist['samples']]
samples_list = self.dataset.remove_duplicates(samples_list)
self.dataset.add(samples_list, 0)
def save_state(self):
"""
Save the process state
"""
if self.output_path is None:
return
logging.info('Save state')
self.dataset.save_dataset(self.output_path / 'dataset.csv')
state = {'history': self.eval.history, 'batch_id': self.batch_id,
'prompt': self.cur_prompt, 'task_description': self.task_description,
'patient': self.patient}
pickle.dump(state, open(self.output_path / 'history.pkl', 'wb'))
def load_state(self, path: str):
"""
Load pretrain state
"""
path = Path(path)
if (path / 'dataset.csv').is_file():
self.dataset.load_dataset(path / 'dataset.csv')
if (path / 'history.pkl').is_file():
state = pickle.load(open(path / 'history.pkl', 'rb'))
self.eval.history = state['history']
self.batch_id = state['batch_id']
self.cur_prompt = state['prompt']
self.task_description = state['task_description']
self.patient = state['patient']
def step(self, current_iter, total_iter):
"""
This is the main optimization process step.
"""
self.log_and_print(f'Starting step {self.batch_id}')
if len(self.dataset.records) == 0:
self.log_and_print('Dataset is empty generating initial samples')
self.generate_initial_samples()
if self.config.use_wandb:
cur_batch = self.dataset.get_leq(self.batch_id)
random_subset = cur_batch.sample(n=min(10, len(cur_batch)))[['text']]
self.wandb_run.log(
{"Prompt": wandb.Html(f"<p>{self.cur_prompt}</p>"), "Samples": wandb.Table(dataframe=random_subset)},
step=self.batch_id)
logging.info('Running annotator')
records = self.annotator.apply(self.dataset, self.batch_id)
self.dataset.update(records)
self.predictor.cur_instruct = self.cur_prompt
logging.info('Running Predictor')
records = self.predictor.apply(self.dataset, self.batch_id, leq=True)
self.dataset.update(records)
self.eval.dataset = self.dataset.get_leq(self.batch_id)
self.eval.eval_score()
logging.info('Calculating Score')
large_errors = self.eval.extract_errors()
self.eval.add_history(self.cur_prompt, self.task_description)
if self.config.use_wandb:
large_errors = large_errors.sample(n=min(6, len(large_errors)))
correct_samples = self.eval.extract_correct()
correct_samples = correct_samples.sample(n=min(6, len(correct_samples)))
vis_data = pd.concat([large_errors, correct_samples])
self.wandb_run.log({"score": self.eval.history[-1]['score'],
"prediction_result": wandb.Table(dataframe=vis_data),
'Total usage': self.calc_usage()}, step=self.batch_id)
if self.stop_criteria():
self.log_and_print('Stop criteria reached')
return True
if current_iter != total_iter-1:
self.run_step_prompt()
self.save_state()
return False
def run_pipeline(self, num_steps: int):
# Run the optimization pipeline for num_steps
num_steps_remaining = num_steps - self.batch_id
for i in range(num_steps_remaining):
stop_criteria = self.step(i, num_steps_remaining)
if stop_criteria:
break
final_result = self.extract_best_prompt()
return final_result