process_output.py (2249B)
1 """Format AEC-t5 output data for evaluate.py. Takes a Path containing .jsonl files as argument.""" 2 import json 3 import argparse 4 import pathlib 5 6 7 def process_output_t5_jsonl_files(directory: pathlib.Path): 8 for filepath in directory.glob("*.jsonl"): 9 # Ignore newly processed outputs 10 if filepath.name.endswith("_evaluate.jsonl"): 11 continue 12 # Declare write path 13 output_path = filepath.with_name(filepath.stem + "_evaluate.test") 14 15 with (filepath.open('r', encoding='utf-8') as infile, output_path.open('w', encoding='utf-8') as outfile): 16 for line in infile: 17 # Read each line containing sample+prediction 18 data = json.loads(line.strip()) 19 # Format prediction dictionary and include sample identifier 20 try: 21 prediction = data["prediction"] 22 prediction["id"] = data["sample"]["id"] 23 del prediction["generated_text"] 24 except KeyError: 25 print(f"{filepath.name} doesn't match the AEC-t5 output format. A sample and a prediction dictionary are expected for each line.") 26 return 27 # Ensure all numeric (deaths) values are converted to integers 28 prediction = {k: int(v) if k.startswith("deaths_") and isinstance(v, str) and v.isdigit() else v for k, v in prediction.items()} 29 for field in prediction: 30 # Force 0 in case of string. 31 # For all runs, strings have consistently returned hallucinations and not potentially correct results such as "5 people" 32 if field.startswith("deaths_") and isinstance(prediction[field], str): 33 prediction[field] = 0 34 json.dump(prediction, outfile) 35 outfile.write("\n") 36 37 38 if __name__ == "__main__": 39 parser = argparse.ArgumentParser(description="Process AEC-t5 output .jsonl files in a given directory.") 40 parser.add_argument("directory", type=str, help="Path to the directory containing AEC-t5 output .jsonl files.") 41 args: argparse.Namespace = parser.parse_args() 42 43 process_output_t5_jsonl_files(pathlib.Path(args.directory))