File size: 5,376 Bytes
42da440 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | import hashlib
import json
import os
import traceback
from augmentoolkit.generation_functions.pipeline_step_class import PipelineStep
# NOTE (performance fix) this stays as it is because it's always just one file.
class SingleGenerationStep(PipelineStep):
def __init__(self, **kwargs):
super().__init__(**kwargs)
def generate_hash(self, input_data):
"""Generate a unique hash from the input data to use as filename"""
# Convert input data to a stable string representation
input_str = json.dumps(input_data, sort_keys=True)
# Create hash
return hashlib.md5(input_str.encode()).hexdigest()
def make_output_path(self, output_dir, input_hash):
"""Returns full path to the output text file based on input hash"""
return os.path.join(output_dir, f"{self.output_file}_{input_hash}.txt")
async def read_previous_output(self, input_data, output_dir):
"""Check if output for this input already exists"""
input_hash = self.generate_hash(input_data)
output_path = self.make_output_path(output_dir, input_hash)
if os.path.exists(output_path):
with open(output_path, "r", encoding="utf-8") as f:
return json.load(f)
return None
def save(
self,
result=None,
input_data=None,
output_dir=None,
full_output=None,
full_response=None,
full_input=None,
include_details=False,
completion_mode=False,
):
"""Save result to a text file named by input hash"""
input_hash = self.generate_hash(input_data)
output_path = self.make_output_path(output_dir, input_hash)
# Ensure directory exists before any file operations
os.makedirs(os.path.dirname(output_path), exist_ok=True)
if self.log_full_outputs and full_output:
print("\n" + "=" * 50)
print(f"FULL OUTPUT FOR HASH {input_hash}:")
print("=" * 50)
# If full_output is a dictionary, print it with nice formatting
if isinstance(full_output, dict):
for key, value in full_output.items():
print(f"\n{key.upper()}:")
print("-" * 50)
print(value)
else:
# Otherwise just print the full output directly
print(full_output)
print("=" * 50 + "\n")
if include_details:
obj = {
self.result_key: result,
"full_output": full_output,
"full_response": full_response,
"full_input": full_input,
"completion_mode": completion_mode,
}
else:
obj = {self.result_key: result}
# Use file lock for safe concurrent writes
# Write to temporary file first then replace for atomic write
temp_path = output_path + ".tmp"
with open(temp_path, "w", encoding="utf-8") as f:
json.dump(obj, f)
os.replace(temp_path, output_path)
return obj
async def run(
self,
input_data=None,
engine_wrapper=None,
default_prompt_folder=None,
prompt_folder=None,
output_dir=None,
completion_mode=False,
use_stop=True,
include_details=False,
**kwargs,
):
full_prompt_path = (
self.prompt_path + ".yaml"
if not completion_mode
else self.prompt_path + ".txt"
)
# Check if this input has already been processed
previous_result = await self.read_previous_output(input_data, output_dir)
if previous_result is not None:
return previous_result
processed_data, additional_kwargs = self.process_input_data(input_data)
error_message = ""
complete = False
max_retries = self.max_retries
while not complete and max_retries > 0:
try:
result, full_output, full_response, full_input = (
await self.generate_data(
processed_data,
engine_wrapper,
full_prompt_path,
prompt_folder,
default_prompt_folder,
completion_mode,
use_stop,
error_message=error_message,
**kwargs,
**additional_kwargs,
)
)
validation_result = self.validation_function(result, input_data)
if validation_result["result"]:
complete = True
else:
error_message = validation_result["message"]
except Exception as e:
print(e)
error_message = str(e)
traceback.print_exc()
max_retries -= 1
if not complete:
return None
return self.save(
result=result,
input_data=input_data,
output_dir=output_dir,
full_output=full_output,
full_response=full_response,
full_input=full_input,
include_details=include_details,
completion_mode=completion_mode,
)
|