diff --git a/requirements.txt b/requirements.txt index 5bdf6448..b4ddf359 100644 --- a/requirements.txt +++ b/requirements.txt @@ -31,4 +31,4 @@ wandb fuzzywuzzy python-Levenshtein anyio>=3.0.0 -grugstream \ No newline at end of file +grugstream>=0.0.13 \ No newline at end of file diff --git a/scripts/ignored_reasoning/james_ignored_reasoning_example.py b/scripts/ignored_reasoning/james_ignored_reasoning_example.py new file mode 100644 index 00000000..94971a27 --- /dev/null +++ b/scripts/ignored_reasoning/james_ignored_reasoning_example.py @@ -0,0 +1,21 @@ +from cot_transparency.formatters.transparency.s1_baselines import ZeroShotCOTUnbiasedTameraTFormatter +from scripts.ignored_reasoning.stage_two import main as stage_two_main +from scripts.ignored_reasoning.stage_two_analysis import aoc_plot, plot_adding_mistakes +from stage_one import main as stage_one_main + +if __name__ == "__main__": + # generate COT + exp_dir = "experiments/james_ignored" + stage_one_main( + tasks=["truthful_qa"], + formatters=[ZeroShotCOTUnbiasedTameraTFormatter.name()], + example_cap=10, + models=["gpt-3.5-turbo", "claude-2"], + temperature=1.0, + exp_dir=exp_dir, + batch=20, + ) + stage_two_dir = "experiments/james_ignored/stage_two" + stage_two_main(input_exp_dir=exp_dir, mistake_model="claude-instant-1", exp_dir=stage_two_dir) + plot_adding_mistakes(exp_dir=stage_two_dir, show_plots=True) + aoc_plot(exp_dir=stage_two_dir, show_plots=True) diff --git a/scripts/ignored_reasoning/stage_two.py b/scripts/ignored_reasoning/stage_two.py index bbf4bd2a..b05adf64 100644 --- a/scripts/ignored_reasoning/stage_two.py +++ b/scripts/ignored_reasoning/stage_two.py @@ -1,8 +1,10 @@ import random from pathlib import Path -from typing import List, Literal, Optional, Type +from typing import List, Literal, NewType, Optional, Type import fire +from git import Sequence +from cot_transparency.apis.base import ModelCaller from cot_transparency.apis.openai.set_key import set_keys_from_env from cot_transparency.data_models.config import CONFIG_MAP @@ -30,7 +32,7 @@ FullCOTFormatter, StageTwoFormatter, ) -from cot_transparency.tasks import run_with_caching_stage_two +from cot_transparency.tasks import run_with_caching_stage_two, task_function from cot_transparency.util import get_exp_dir_name from stage_one import get_valid_stage1_formatters @@ -112,68 +114,88 @@ def get_early_answering_tasks( return outputs +def create_mistake_task_spec_for_stage_one( + stage_one_output: TaskOutput, + exp_dir: str, + mistake_adding_model: str, + mistake_adding_temperature: float, + n_mistake_insertion_points: int, +) -> list[StageTwoTaskSpec]: + DataExample = task_name_to_data_example(stage_one_output.task_spec.task_name) + data_example = stage_one_output.task_spec.read_data_example_or_raise(DataExample) + original_question: str = data_example.get_parsed_input() + + cot_steps = get_cot_steps(stage_one_output.first_raw_response) + cot_steps = filter_cot_by_possible_ends(cot_steps) + + rng = random.Random(original_question) + # we don't want to insert mistakes at the first step, because then there is no sentence to make a mistake in + # as the first one is blank cot + if len(cot_steps) == 0: + print("WARNING - skipping task as len(cot_steps) == 0") + return [] + + sample_idxs = [ + 0, + len(cot_steps) - 1, + ] # we always do the first and last one to get good aoc numbers + if len(cot_steps) > 2: + mistake_idxs = rng.sample( + range(1, len(cot_steps) - 1), + min(n_mistake_insertion_points, len(cot_steps) - 2), + ) + sample_idxs.extend(mistake_idxs) + sample_idxs = sorted(list(set(sample_idxs))) + + config = CONFIG_MAP[mistake_adding_model].copy() + config.max_tokens = 100 + config.temperature = mistake_adding_temperature + config.stop = ["\n", "```"] + + original_model_that_generated_cot = stage_one_output.task_spec.inference_config.model + path = Path( + f"{exp_dir}/mistakes_stage1/s1-{stage_one_output.task_spec.formatter_name}/cot-{original_model_that_generated_cot}/{stage_one_output.task_spec.task_name}/" + f"/mistake-{config.model}/{FewShotGenerateMistakeFormatter.name()}.json" + ) + out: list[StageTwoTaskSpec] = [] + for i in sample_idxs: + messages = FewShotGenerateMistakeFormatter.format_example( + original_question=original_question, sentence=cot_steps[i] + ) + task_spec = StageTwoTaskSpec( + stage_one_output=stage_one_output, + inference_config=config, + formatter_name=FewShotGenerateMistakeFormatter.name(), + messages=messages, + out_file_path=path, + trace_info=TraceInfo(original_cot=cot_steps, mistake_inserted_idx=i), + ) + out.append(task_spec) + return out + + def get_mistakes( - stage_one_outputs: list[TaskOutput], + stage_one_outputs: Sequence[TaskOutput], exp_dir: str, batch: int = 10, mistake_adding_model: str = "text-davinci-002", mistake_adding_temperature: float = 1.0, save_mistake_generating_file_every: int = 50, n_mistake_insertion_points: int = 3, -) -> list[StageTwoTaskOutput]: +) -> Sequence[StageTwoTaskOutput]: # mistakes we need to make call to api to generate the mistake, we may use a different model here # e.g Tamera et al use a non RLHF model to generate mistakes specs: list[StageTwoTaskSpec] = [] for stage_one_output in stage_one_outputs: - DataExample = task_name_to_data_example(stage_one_output.task_spec.task_name) - data_example = stage_one_output.task_spec.read_data_example_or_raise(DataExample) - original_question: str = data_example.get_parsed_input() - - cot_steps = get_cot_steps(stage_one_output.first_raw_response) - cot_steps = filter_cot_by_possible_ends(cot_steps) - - rng = random.Random(original_question) - # we don't want to insert mistakes at the first step, because then there is no sentence to make a mistake in - # as the first one is blank cot - if len(cot_steps) == 0: - print("WARNING - skipping task as len(cot_steps) == 0") - continue - - sample_idxs = [ - 0, - len(cot_steps) - 1, - ] # we always do the first and last one to get good aoc numbers - if len(cot_steps) > 2: - mistake_idxs = rng.sample( - range(1, len(cot_steps) - 1), - min(n_mistake_insertion_points, len(cot_steps) - 2), - ) - sample_idxs.extend(mistake_idxs) - sample_idxs = sorted(list(set(sample_idxs))) - - config = CONFIG_MAP[mistake_adding_model].copy() - config.max_tokens = 100 - config.temperature = mistake_adding_temperature - config.stop = ["\n", "```"] - - original_model_that_generated_cot = stage_one_output.task_spec.inference_config.model - path = Path( - f"{exp_dir}/mistakes_stage1/s1-{stage_one_output.task_spec.formatter_name}/cot-{original_model_that_generated_cot}/{stage_one_output.task_spec.task_name}/" - f"/mistake-{config.model}/{FewShotGenerateMistakeFormatter.name()}.json" - ) - for i in sample_idxs: - messages = FewShotGenerateMistakeFormatter.format_example( - original_question=original_question, sentence=cot_steps[i] - ) - task_spec = StageTwoTaskSpec( + specs.extend( + create_mistake_task_spec_for_stage_one( stage_one_output=stage_one_output, - inference_config=config, - formatter_name=FewShotGenerateMistakeFormatter.name(), - messages=messages, - out_file_path=path, - trace_info=TraceInfo(original_cot=cot_steps, mistake_inserted_idx=i), + exp_dir=exp_dir, + mistake_adding_model=mistake_adding_model, + mistake_adding_temperature=mistake_adding_temperature, + n_mistake_insertion_points=n_mistake_insertion_points, ) - specs.append(task_spec) + ) print(f"1. Generating mistakes using {mistake_adding_model}") print("n specs", len(specs)) @@ -194,6 +216,10 @@ def get_mistakes( outputs = run_with_caching_stage_two(save_mistake_generating_file_every, batch, specs) + return filter_mistakes_output(outputs) + + +def filter_mistakes_output(outputs: Sequence[StageTwoTaskOutput]) -> list[StageTwoTaskOutput]: # Discard outputs where there was no output filtered_outputs = [output for output in outputs if output.first_parsed_response is not None] n_discarded = len(outputs) - len(filtered_outputs) @@ -204,17 +230,75 @@ def get_mistakes( filtered_outputs = [i for i in outputs if "NO_REASONING" not in i.first_parsed_response] # type: ignore n_discarded = len(outputs) - len(filtered_outputs) print(f"Discarding {n_discarded} outputs where the mistake model didn't deem there to be a reasoning step") - outputs = filtered_outputs + return filtered_outputs + + +# A Newtype so that you won't get confused with all the different types +RecomputeTaskSpec = NewType("RecomputeTaskSpec", StageTwoTaskSpec) + + +def mistakes_into_completed_cot_spec( + mistake: StageTwoTaskOutput, + exp_dir: str, +) -> Optional[RecomputeTaskSpec]: + generated_mistake = mistake + if generated_mistake.first_parsed_response is None or "NO_REASONING" in generated_mistake.first_parsed_response: + print("WARNING - skipping task as NO_REASONING found in the trace passed to the mistake generator") + return None + else: + stage_one_output = generated_mistake.task_spec.stage_one_output + config = stage_one_output.task_spec.inference_config.copy() + + path = Path( + f"{exp_dir}/mistakes_stage2/s1-{stage_one_output.task_spec.formatter_name}/{stage_one_output.task_spec.task_name}/{config.model}/{CompletePartialCOT.name()}.json" + ) + + trace_info: TraceInfo | None = generated_mistake.task_spec.trace_info + assert trace_info is not None + trace_info = trace_info.model_copy() + trace_info.sentence_with_mistake = generated_mistake.first_parsed_response + + partial_cot_trace = trace_info.get_trace_upto_mistake() + + messages = CompletePartialCOT.format_example( + stage_one_output.task_spec.messages, partial_cot_trace, config.model + ) + + task_spec = StageTwoTaskSpec( + stage_one_output=stage_one_output, + inference_config=config, + formatter_name=CompletePartialCOT.name(), + messages=messages, + out_file_path=path, + trace_info=trace_info, + ) + + return RecomputeTaskSpec(task_spec) - return outputs + +def execute_recomputation(task_spec: RecomputeTaskSpec, caller: ModelCaller) -> Sequence[StageTwoTaskOutput]: + trace_info = task_spec.trace_info + assert trace_info + # if the mistake was the last step in the reasoning trace, then we don't need to complete the COT + # so just make a task output with no response + if trace_info.mistake_inserted_idx == len(trace_info.original_cot) - 1: + output = StageTwoTaskOutput( + task_spec=task_spec, + inference_output=ModelOutput(raw_response="", parsed_response=""), + ) + return [output for _ in range(task_spec.inference_config.n)] + else: + # We need to complete the COT + return task_function(task=task_spec, raise_after_retries=False, caller=caller) def recomplete_cot_with_inserted_mistake( - generated_mistakes: list[StageTwoTaskOutput], + generated_mistakes: Sequence[StageTwoTaskOutput], exp_dir: str, save_completing_with_mistakes_every: int = 50, batch: int = 10, ) -> List[StageTwoTaskOutput]: + """Note: This is an old function that isn't used anymore by the stage_two_stream.py script""" mistakes_inserted_at_last_position: list[StageTwoTaskOutput] = [] specs: list[StageTwoTaskSpec] = [] @@ -266,6 +350,52 @@ def recomplete_cot_with_inserted_mistake( return outputs + mistakes_inserted_at_last_position +def single_get_best_single_answer_tasks_given_mistakes( + cot_with_mistakes_outputs: StageTwoTaskOutput, + exp_dir: str, + temperature: Optional[float] = None, +) -> Optional[StageTwoTaskSpec]: + output = cot_with_mistakes_outputs + stage_one_output = output.task_spec.stage_one_output + config = stage_one_output.task_spec.inference_config.copy() + config.max_tokens = 30 # code-davinci-002 doesn't return answer unless we set this to greater than 1 + if temperature is not None: + config.temperature = temperature + + parsed_response: str | None = output.first_parsed_response + if parsed_response is None: + print("WARNING - skipping task as parsed_response is None") + return None + + if stage_one_output.task_spec.formatter_name == FewShotCOTUnbiasedCompletionNoRoleTameraTFormatter.name(): + Formatter = FullCOTCompletionFormatter + else: + Formatter = FullCOTFormatter + + path = Path( + f"{exp_dir}/mistakes_final/s1-{stage_one_output.task_spec.formatter_name}/{stage_one_output.task_spec.task_name}/{config.model}/{Formatter.name()}.json" + ) + trace_info = output.task_spec.trace_info + assert trace_info is not None + trace_info.regenerated_cot_post_mistake = parsed_response + cot_trace_with_mistake = trace_info.get_complete_modified_cot() + + final_task = StageTwoTaskSpec( + stage_one_output=output.task_spec.stage_one_output, + inference_config=config, + formatter_name=Formatter.name(), + messages=Formatter.format_example( + stage_one_output.task_spec.messages, + cot_trace_with_mistake, + config.model, + ), + out_file_path=path, + n_steps_in_cot_trace=len(get_cot_steps(cot_trace_with_mistake)), + trace_info=trace_info, + ) + return final_task + + def get_best_single_answer_tasks_given_mistakes( cots_with_mistakes_outputs: list[StageTwoTaskOutput], exp_dir: str, @@ -273,44 +403,11 @@ def get_best_single_answer_tasks_given_mistakes( ) -> list[StageTwoTaskSpec]: specs: List[StageTwoTaskSpec] = [] for output in cots_with_mistakes_outputs: - stage_one_output = output.task_spec.stage_one_output - config = stage_one_output.task_spec.inference_config.copy() - config.max_tokens = 30 # code-davinci-002 doesn't return answer unless we set this to greater than 1 - if temperature is not None: - config.temperature = temperature - - parsed_response: str | None = output.first_parsed_response - if parsed_response is None: - print("WARNING - skipping task as parsed_response is None") - continue - - if stage_one_output.task_spec.formatter_name == FewShotCOTUnbiasedCompletionNoRoleTameraTFormatter.name(): - Formatter = FullCOTCompletionFormatter - else: - Formatter = FullCOTFormatter - - path = Path( - f"{exp_dir}/mistakes_final/s1-{stage_one_output.task_spec.formatter_name}/{stage_one_output.task_spec.task_name}/{config.model}/{Formatter.name()}.json" - ) - trace_info = output.task_spec.trace_info - assert trace_info is not None - trace_info.regenerated_cot_post_mistake = parsed_response - cot_trace_with_mistake = trace_info.get_complete_modified_cot() - - final_task = StageTwoTaskSpec( - stage_one_output=output.task_spec.stage_one_output, - inference_config=config, - formatter_name=Formatter.name(), - messages=Formatter.format_example( - stage_one_output.task_spec.messages, - cot_trace_with_mistake, - config.model, - ), - out_file_path=path, - n_steps_in_cot_trace=len(get_cot_steps(cot_trace_with_mistake)), - trace_info=trace_info, + final_task = single_get_best_single_answer_tasks_given_mistakes( + cot_with_mistakes_outputs=output, exp_dir=exp_dir, temperature=temperature ) - specs.append(final_task) + if final_task: + specs.append(final_task) return specs diff --git a/scripts/ignored_reasoning/stage_two_analysis.py b/scripts/ignored_reasoning/stage_two_analysis.py index 92b1cad6..0068fffe 100644 --- a/scripts/ignored_reasoning/stage_two_analysis.py +++ b/scripts/ignored_reasoning/stage_two_analysis.py @@ -4,11 +4,15 @@ import numpy as np import pandas as pd import seaborn as sns +from git import Sequence from matplotlib import pyplot as plt from analysis import TASK_MAP, accuracy_for_df, get_general_metrics from cot_transparency.data_models.io import ExpLoader -from cot_transparency.data_models.models import StageTwoExperimentJsonFormat, TaskOutput +from cot_transparency.data_models.models import ( + StageTwoTaskOutput, +) +from cot_transparency.data_models.models import TaskOutput from cot_transparency.formatters.transparency.trace_manipulation import get_cot_steps # Used to produce human readable names on plots @@ -52,6 +56,19 @@ def get_aoc(df: pd.DataFrame, x="cot_trace_length") -> pd.DataFrame: def get_auc(group: pd.DataFrame) -> float: assert group["original_cot_trace_length"].nunique() == 1 proportion_of_cot = group[x] / max(group[x]) + """ + TODO: Maybe we need to renormalize the proportion_of_cot to be between 0 and 1? + Right now it can be between e.g. 0.6-1.0, which messes up np.trapz? + And sometimes its one single point too, which also messes up np.trapz + min_ = min(proportion_of_cot) + max_ = max(proportion_of_cot) + try: + proportion_scaled = proportion_of_cot.apply(lambda x: (x - min_) / (max_ - min_)) + except Exception as e: + print("Maybe there is only one value so min_ == max_") + raise e + """ + auc = np.trapz(group[0], x=proportion_of_cot) return auc @@ -76,11 +93,9 @@ def get_auc(group: pd.DataFrame) -> float: return areas -def convert_stage2_experiment_to_dataframe( - exp: StageTwoExperimentJsonFormat, -) -> pd.DataFrame: +def convert_stage2_experiment_to_dataframe(exp: Sequence[StageTwoTaskOutput]) -> pd.DataFrame: out = [] - for task_output in exp.outputs: + for task_output in exp: d_with_config = get_general_metrics(task_output) d_with_config["model"] = task_output.task_spec.inference_config.model d_with_config["task_name"] = task_output.task_spec.stage_one_output.task_spec.task_name @@ -115,7 +130,7 @@ def get_data_frame_from_exp_dir(exp_dir: str) -> pd.DataFrame: loaded_dict = ExpLoader.stage_two(exp_dir, final_only=True) dfs = [] for exp in loaded_dict.values(): - df = convert_stage2_experiment_to_dataframe(exp) + df = convert_stage2_experiment_to_dataframe(exp.outputs) dfs.append(df) df = pd.concat(dfs) df["is_correct"] = (df.parsed_response == df.ground_truth).astype(int) @@ -126,6 +141,16 @@ def get_data_frame_from_exp_dir(exp_dir: str) -> pd.DataFrame: return df # type: ignore +def get_data_frame_from_exp_dir_items(items: Sequence[StageTwoTaskOutput]) -> pd.DataFrame: + df = convert_stage2_experiment_to_dataframe(items) + df["is_correct"] = (df.parsed_response == df.ground_truth).astype(int) + # filter out the NOT_FOUND rows + n_not_found = len(df[df.parsed_response == "NOT_FOUND"]) + print(f"Number of NOT_FOUND rows: {n_not_found}") + df = df[df.parsed_response != "NOT_FOUND"] + return df # type: ignore + + def plot_historgram_of_lengths( exp_dir: str, ): @@ -151,6 +176,31 @@ def plot_historgram_of_lengths( plt.show() +def plot_histogram_from_list( + items: Sequence[StageTwoTaskOutput], +): + df = get_data_frame_from_exp_dir_items(items) + + hue = "task_name" + x = "CoT Length" + col = "model" + y = "Counts" + + # rename "original_cot_trace_length" to "CoT Length" + df = df.rename(columns={"original_cot_trace_length": x}) + + # for histogram we want the counts of the original_cot_trace_length + # filter on the unique stage_one_hash + counts = df.groupby([hue, col, x]).stage_one_hash.nunique().reset_index() + counts = counts.rename(columns={"stage_one_hash": y}) + + # facet plot of the proportion of the trace, break down by original_cot_trace_length + g = sns.FacetGrid(counts, col=col, col_wrap=2, legend_out=True) + g.map_dataframe(sns.barplot, x=x, y=y, hue=hue) + g.add_legend() + plt.show() + + def df_filters( df: pd.DataFrame, inconsistent_only: bool, @@ -270,7 +320,35 @@ def plot_early_answering( col: str = "original_cot_trace_length", hue: str = "model", ): - df = get_data_frame_from_exp_dir(exp_dir) + items: list[StageTwoTaskOutput] = [] + loaded_dict = ExpLoader.stage_two(exp_dir, final_only=True) + for vals in loaded_dict.values(): + outputs = vals.outputs + items.extend(outputs) + + return plot_early_answering_from_list( + items=items, + show_plots=show_plots, + inconsistent_only=inconsistent_only, + aggregate_over_tasks=aggregate_over_tasks, + model_filter=model_filter, + length_filter=length_filter, + col=col, + hue=hue, + ) + + +def plot_early_answering_from_list( + items: Sequence[StageTwoTaskOutput], + show_plots: bool = False, + inconsistent_only: bool = False, + aggregate_over_tasks: bool = False, + model_filter: Optional[str] = None, + length_filter: Optional[list[int]] = None, + col: str = "original_cot_trace_length", + hue: str = "model", +): + df = get_data_frame_from_exp_dir_items(items=items) # drop formatters that have Mistake in the name df = df[~df.has_mistake] df = df_filters(df, inconsistent_only, aggregate_over_tasks, model_filter, length_filter) # type: ignore @@ -303,7 +381,35 @@ def plot_adding_mistakes( col: str = "original_cot_trace_length", hue: str = "model", ): - df = get_data_frame_from_exp_dir(exp_dir) + items: list[StageTwoTaskOutput] = [] + loaded_dict = ExpLoader.stage_two(exp_dir, final_only=True) + for vals in loaded_dict.values(): + outputs = vals.outputs + items.extend(outputs) + + return plot_adding_mistakes_from_list( + items=items, + show_plots=show_plots, + inconsistent_only=inconsistent_only, + aggregate_over_tasks=aggregate_over_tasks, + model_filter=model_filter, + length_filter=length_filter, + col=col, + hue=hue, + ) + + +def plot_adding_mistakes_from_list( + items: Sequence[StageTwoTaskOutput], + show_plots: bool = False, + inconsistent_only: bool = False, + aggregate_over_tasks: bool = False, + model_filter: Optional[str] = None, + length_filter: Optional[list[int]] = None, + col: str = "original_cot_trace_length", + hue: str = "model", +): + df = convert_stage2_experiment_to_dataframe(items) df = df[~df.was_truncated] df = df_filters(df, inconsistent_only, aggregate_over_tasks, model_filter, length_filter) # type: ignore @@ -333,7 +439,33 @@ def aoc_plot( length_filter: Optional[list[int]] = None, hue: str = "stage_one_formatter_name", ): - df = get_data_frame_from_exp_dir(exp_dir) + items: list[StageTwoTaskOutput] = [] + loaded_dict = ExpLoader.stage_two(exp_dir, final_only=True) + for vals in loaded_dict.values(): + outputs = vals.outputs + items.extend(outputs) + + return aoc_plot_from_list( + items=items, + show_plots=show_plots, + inconsistent_only=inconsistent_only, + aggregate_over_tasks=aggregate_over_tasks, + model_filter=model_filter, + length_filter=length_filter, + hue=hue, + ) + + +def aoc_plot_from_list( + items: Sequence[StageTwoTaskOutput], + show_plots: bool = False, + inconsistent_only: bool = False, + aggregate_over_tasks: bool = False, + model_filter: Optional[str] = None, + length_filter: Optional[list[int]] = None, + hue: str = "stage_one_formatter_name", +): + df = get_data_frame_from_exp_dir_items(items) df = df_filters(df, inconsistent_only, aggregate_over_tasks, model_filter, length_filter) # Mistakes AoC diff --git a/scripts/ignored_reasoning/stage_two_stream.py b/scripts/ignored_reasoning/stage_two_stream.py new file mode 100644 index 00000000..54d1a55b --- /dev/null +++ b/scripts/ignored_reasoning/stage_two_stream.py @@ -0,0 +1,153 @@ +import asyncio +from pathlib import Path + +from anyio import CapacityLimiter + +from grugstream import Observable + +from cot_transparency.apis import UniversalCaller +from cot_transparency.data_models.models import StageTwoTaskOutput +from cot_transparency.formatters.transparency.s1_baselines import ZeroShotCOTUnbiasedTameraTFormatter +from cot_transparency.tasks import task_function +from scripts.ignored_reasoning.stage_two import ( + create_mistake_task_spec_for_stage_one, + execute_recomputation, + get_early_answering_tasks, + mistakes_into_completed_cot_spec, + single_get_best_single_answer_tasks_given_mistakes, +) +from scripts.ignored_reasoning.stage_two_analysis import ( + aoc_plot_from_list, + plot_early_answering_from_list, + plot_histogram_from_list, + plot_adding_mistakes_from_list, +) +from cot_transparency.streaming.stage_one_stream import stage_one_stream + + +async def main(): + stage_one_path = Path("experiments/stream_mistakes/stage_one.jsonl") + stage_one_caller = UniversalCaller().with_file_cache(stage_one_path) + + recompute_cot_caller = UniversalCaller().with_file_cache("experiments/stream_mistakes/recompute_cot.jsonl") + add_mistake_caller = UniversalCaller().with_file_cache("experiments/stream_mistakes/add_mistakes.jsonl") + final_answer_caller = UniversalCaller().with_file_cache("experiments/stream_mistakes/final_answer.jsonl") + n_mistake_insertion_points = 16 + stage_one_obs = stage_one_stream( + formatters=[ZeroShotCOTUnbiasedTameraTFormatter.name()], + # hacked truthful_qa to have correct labels as A + tasks=["truthful_qa"], + example_cap=200, + raise_after_retries=False, + temperature=1.0, + caller=stage_one_caller, + batch=20, + models=["gpt-3.5-turbo", "claude-2"], + ) + # shared capacity limit between the stages + capacity_limit = CapacityLimiter(50) + + early_answer_obs = ( + stage_one_obs.map( + lambda task_output: get_early_answering_tasks( + stage_one_output=task_output, + exp_dir="not_used", + temperature=None, + n_samples_per_cot=4, + full_answers_only=False, + ) + ) + .flatten_list() + .map_blocking_par( + lambda stage_two_spec: task_function( + task=stage_two_spec, raise_after_retries=False, caller=recompute_cot_caller + ) + ) + .flatten_list() + ) + early_answer_results = await early_answer_obs.to_list() + plot_early_answering_from_list(items=early_answer_results, show_plots=True) + + mistakes_obs: Observable[StageTwoTaskOutput] = ( + stage_one_obs.map( + # Create mistake spec + lambda task_output: create_mistake_task_spec_for_stage_one( + stage_one_output=task_output, + exp_dir="not_used", + mistake_adding_temperature=1.0, + n_mistake_insertion_points=n_mistake_insertion_points, + mistake_adding_model="claude-instant-1", + ) + ) + .flatten_list() + # Call the mistake making model + .map_blocking_par( + lambda stage_two_spec: task_function( + task=stage_two_spec, raise_after_retries=False, caller=add_mistake_caller + ), + max_par=capacity_limit, + ) + .flatten_list() + # We want only not None responses + .filter(lambda task: task.first_parsed_response is not None) + # # Make another spec! + .map(lambda output: mistakes_into_completed_cot_spec(mistake=output, exp_dir="not_used")) + .flatten_optional() + # Execute recomputation + .map_blocking_par( + lambda spec: execute_recomputation(task_spec=spec, caller=recompute_cot_caller), max_par=capacity_limit + ) + .flatten_list() + # final best answer task spec + .map( + lambda x: single_get_best_single_answer_tasks_given_mistakes( + cot_with_mistakes_outputs=x, exp_dir="not_used" + ) ## prev up to here + ) + .flatten_optional() + .map_blocking_par( + lambda stage_two_spec: task_function( + task=stage_two_spec, raise_after_retries=False, caller=final_answer_caller + ), + max_par=capacity_limit, + ) + .flatten_list() + .tqdm(None) + ) + + baseline_no_mistakes = ( + stage_one_obs.map( + lambda stage_one_task: get_early_answering_tasks( + stage_one_output=stage_one_task, + exp_dir="not_used", + temperature=stage_one_task.task_spec.inference_config.temperature, + full_answers_only=True, + ) + ) + .flatten_list() + .map_blocking_par( + lambda stage_two_spec: task_function( + task=stage_two_spec, raise_after_retries=False, caller=final_answer_caller + ), + max_par=capacity_limit, + ) + .flatten_list() + ) + + mistakes_results = await mistakes_obs.to_list() + baseline_no_mistakes_results = await baseline_no_mistakes.to_slist() + print(f"length baseline_no_mistakes_results { baseline_no_mistakes_results.length}") + all_ = mistakes_results + baseline_no_mistakes_results + print("done with mistakes") + plot_histogram_from_list(all_) + aoc_plot_from_list(all_, show_plots=True) + plot_adding_mistakes_from_list(mistakes_results + baseline_no_mistakes_results, show_plots=True) + + stage_one_caller.save_cache() + recompute_cot_caller.save_cache() + add_mistake_caller.save_cache() + final_answer_caller.save_cache() + + +if __name__ == "__main__": + asyncio.run(main())