import json import os from typing import Optional from dataclasses import dataclass, field import numpy as np import torch import blobfile as bf from summarize_from_feedback import task_data from summarize_from_feedback.datasets import jsonl_encoding from summarize_from_feedback.utils import exact_div, Timer, hyperparams from summarize_from_feedback.utils.logging_utils import setup_logging_with_pacific_tz from summarize_from_feedback.utils.torch_utils import label_logprobs from summarize_from_feedback.utils.nested import map_nested from summarize_from_feedback.utils.assertions import assert_shape_eq, assert_eq from summarize_from_feedback.query_response_model import ModelSpec, SampleHParams, PADDING_TOKEN from summarize_from_feedback.policy import Policy from summarize_from_feedback.tasks import TaskHParams, ResponseEncoder @dataclass class HParams(hyperparams.HParams): model_spec: ModelSpec = field(default_factory=ModelSpec) orig_model_spec: Optional[ModelSpec] = None task: TaskHParams = field(default_factory=TaskHParams) query_dataset_split: str = None sample: SampleHParams = field(default_factory=SampleHParams) num_queries: int = None # Note that these batch sizes are in # of datapoints; each datapoint may turn into multiple # sequences fed to the model queries_per_run_per_replica: int = 1 responses_per_query: int = 1 responses_per_query_per_batch: int = 1 seed: int = 0 fp16_activations: bool = True INVALID_LOGPROB = 1.0 def avg_negative(x): mask = x <= 0 return np.sum(x * mask) / np.sum(mask) def main(H: HParams): layout = H.model_spec.run_params.all_gpu_layout() # Instantiate policy policy = Policy(task_hparams=H.task, spec=H.model_spec, layout=layout) if H.orig_model_spec: assert H.orig_model_spec.run_params.n_shards == H.model_spec.run_params.n_shards orig_policy = Policy(task_hparams=H.task, spec=H.orig_model_spec, layout=layout) else: orig_policy = None encoder = policy.encoder response_encoder = ResponseEncoder(H.task.response, encoder) setup_logging_with_pacific_tz() act_dtype = torch.float16 if H.fp16_activations else torch.float32 is_logging_rank = layout.is_logging_rank total_queries_per_replica = exact_div(H.num_queries, layout.n_replicas) num_runs = exact_div(total_queries_per_replica, H.queries_per_run_per_replica) input_iter = task_data.get_iter_for_task( H.task, encoder=encoder, dataset_split=H.query_dataset_split, batch_size=H.queries_per_run_per_replica, layout=layout, seed=H.seed, all_fields=True, ) log_dir = os.getenv("OUTPUT_DIR") or os.path.join("/tmp/jobs", os.getenv("JOB_NAME")) results_dir = os.path.join(log_dir, "results") bf.makedirs(results_dir) with open(os.path.join(log_dir, "hparams.json"), "w") as f: json.dump(H.to_json(), f, indent=2) if is_logging_rank: with open(os.path.join(results_dir, "task_hparams.json"), "w") as f: json.dump(H.task.to_json(), f) with open(os.path.join(results_dir, "hparams.json"), "w") as f: json.dump(H.to_json(), f) # Creates files for printing. Only the replica root prints the files local_file_name = os.devnull if layout.is_replica_root: fname = f"samples.{layout.replica_idx}.jsonl" local_file_name = os.path.join(results_dir, fname) print(f"Samples will be written to {local_file_name}") def prepare_eval_fn_and_inputs(tokens): def eval_fn(outputs_mb, eval_inputs_mb): logprobs = label_logprobs(logits=outputs_mb["logits"], labels=eval_inputs_mb["labels"]) logprobs = torch.masked_fill(logprobs, eval_inputs_mb["mask"], INVALID_LOGPROB) return dict(logprobs=logprobs) mask = tokens == PADDING_TOKEN return eval_fn, dict(labels=torch.masked_fill(tokens, mask, 0), mask=mask) runs_per_query = exact_div(H.responses_per_query, H.responses_per_query_per_batch) with open(local_file_name, "w") as f: for run_idx in range(num_runs): with Timer() as timer: input = next(input_iter) context_tokens = input["context"]["tokens"] assert_shape_eq( context_tokens, (H.queries_per_run_per_replica, H.task.query.length), "Context tokens shape mismatch", ) ref_tokens = input["reference"]["tokens"].unsqueeze(1) assert_shape_eq( ref_tokens, (H.queries_per_run_per_replica, 1, H.task.response.length), "Ref tokens shape mismatch", ) # Sample from policy all_sample_results = [] for _ in range(runs_per_query): sample_results = policy.sample( context_tokens, responses_per_query=H.responses_per_query_per_batch, sample_H=H.sample, act_dtype=act_dtype, ) assert_shape_eq( sample_results["samples"], ( H.queries_per_run_per_replica, H.responses_per_query_per_batch, H.task.response.length, ), "Samples size mismatch", ) processed_samples = response_encoder.process_responses( sample_results["samples"] ) sample_results["processed_samples"] = processed_samples assert_shape_eq( processed_samples, ( H.queries_per_run_per_replica, H.responses_per_query_per_batch, H.task.response.length, ), "Samples size mismatch", ) sample_results["logprobs"] = torch.masked_fill( sample_results["logprobs"], processed_samples == PADDING_TOKEN, INVALID_LOGPROB, ) if orig_policy is not None: eval_fn, eval_inputs = prepare_eval_fn_and_inputs(processed_samples) orig_eval_results = orig_policy.eval( context_tokens, processed_samples, eval_fn=eval_fn, eval_inputs=eval_inputs, act_dtype=act_dtype, ) sample_results["orig_eval_results"] = orig_eval_results sample_results = map_nested(sample_results, lambda x: x.cpu().numpy()) all_sample_results.append(sample_results) eval_fn, eval_inputs = prepare_eval_fn_and_inputs(ref_tokens) ref_eval_results = policy.eval( context_tokens, ref_tokens, eval_fn=eval_fn, eval_inputs=eval_inputs, act_dtype=act_dtype, ) if orig_policy is not None: orig_ref_eval_results = orig_policy.eval( context_tokens, ref_tokens, eval_fn=eval_fn, eval_inputs=eval_inputs, act_dtype=act_dtype, ) if layout.is_replica_root: for batch_idx in range(H.queries_per_run_per_replica): context_tokens = sample_results["contexts"][batch_idx] context = encoder.decode(context_tokens) # Dump to a file so that we can use things in downstream tasks # samples (written to file) is now a list of strings d = dict(context=context, context_tokens=context_tokens) d["sample_tokens"] = np.concatenate( [ sample_results["processed_samples"][batch_idx] for sample_results in all_sample_results ], axis=0, ) assert_shape_eq( d["sample_tokens"], (H.responses_per_query, H.task.response.length), "Sample tokens shape mismatch", ) d["samples"] = response_encoder.decode_responses(d["sample_tokens"]) d["logprobs"] = np.concatenate( [ sample_results["logprobs"][batch_idx] for sample_results in all_sample_results ], axis=0, ) assert_shape_eq( d["logprobs"], (H.responses_per_query, H.task.response.length), "Logprobs shape mismatch", ) if orig_policy is not None: d["orig_logprobs"] = np.concatenate( [ sample_results["orig_eval_results"]["eval_stats"]["logprobs"][ batch_idx ] for sample_results in all_sample_results ], axis=0, ) assert_shape_eq( d["orig_logprobs"], (H.responses_per_query, H.task.response.length), "Orig logprobs shape mismatch", ) # Process ref_tokens (H.queries_per_run_per_replica, H.task.response.length) d["ref_tokens"] = ref_tokens[batch_idx].squeeze(0).cpu().numpy() d["ref"] = response_encoder.decode_response(d["ref_tokens"]) assert_eq( len(d["ref_tokens"]), H.task.response.length, "Ref tokens shape mismatch", ) d["ref_logprobs"] = ( ref_eval_results["eval_stats"]["logprobs"][batch_idx] .squeeze(0) .cpu() .numpy() ) assert_eq( len(d["ref_logprobs"]), H.task.response.length, "Ref logprobs shape mismatch", ) if orig_policy is not None: d["orig_ref_logprobs"] = ( orig_ref_eval_results["eval_stats"]["logprobs"][batch_idx] .squeeze(0) .cpu() .numpy() ) assert_eq( len(d["orig_ref_logprobs"]), H.task.response.length, "Orig ref Logprobs shape mismatch", ) if "extra_fields" in input: d["extra_fields"] = input["extra_fields"][batch_idx] print("=" * 80) replica_sample_idx = run_idx * H.queries_per_run_per_replica + batch_idx print(f"RESULT {replica_sample_idx} of {total_queries_per_replica}") print(f"CONTEXT:") print(context) print(f"REF:") print(d["ref"]) print("avg logprob", avg_negative(d["ref_logprobs"])) if orig_policy is not None: print("avg orig logprob", avg_negative(d["orig_ref_logprobs"])) for sample_idx in range(H.responses_per_query): print(f"SAMPLE {sample_idx}:") print(d["samples"][sample_idx]) print("avg logprob", avg_negative(d["logprobs"][sample_idx])) if orig_policy is not None: print( "avg orig logprob", avg_negative(d["orig_logprobs"][sample_idx]) ) f.write((json.dumps(jsonl_encoding.encode_example(d)) + "\n")) if layout.is_replica_root: print(f"Batch {run_idx+1} of {num_runs}. Took {timer.interval} seconds") return dict(output_path=results_dir)