kernelfoundry.algorithm.utils.extract_code

Functions

apply_diff(original_code, diff_text)	Apply a diff to the original code
extract_code_blocks(text, code_language_types)	Extract all code blocks from text, combine them to return as a single string
extract_code_flexible(output_string[, tag, ...])	Extracts code either from codeblocks, from string, or from tags
extract_code_from_tags(output_string[, tag, ...])
extract_cpp_code_heuristic(output_string)	Extract C++ code from model output
extract_diffs(diff_text)	Extract diff blocks from the diff text
extract_first_code(output_string, ...)	Extract first code block from model output, specified by code_language_type
extract_last_code(output_string, ...)	Extract last code block from model output, specified by code_language_type
extract_python_code(text)	Extract python code from model output
format_diff_summary(diff_blocks)	Create a human-readable summary of the diff
remove_code_block_header(code, ...)	Assume input is code but just with like python, cpp, etc.
replace_function_calls(text)	For templated kernels, add template_args to pytorch functional

kernelfoundry.algorithm.utils.extract_code.extract_python_code(text)

Extract python code from model output

kernelfoundry.algorithm.utils.extract_code.remove_code_block_header(code, code_language_type)

Assume input is code but just with like python, cpp, etc. at the top

kernelfoundry.algorithm.utils.extract_code.extract_first_code(output_string: str, code_language_types: list[str]) → str

Extract first code block from model output, specified by code_language_type

kernelfoundry.algorithm.utils.extract_code.extract_code_flexible(output_string: str, tag: str = 'cuda', use_first: bool = False, code_language_types: list = ['cuda', 'sycl', 'python', 'cpp', 'triton', 'ocl']) → str | None

Extracts code either from codeblocks, from string, or from tags

kernelfoundry.algorithm.utils.extract_code.extract_cpp_code_heuristic(output_string: str) → str | None

Extract C++ code from model output

kernelfoundry.algorithm.utils.extract_code.extract_code_from_tags(output_string: str, tag: str = 'cuda', use_first: bool = False, code_language_types: list = ['cuda', 'sycl', 'python', 'cpp', 'triton', 'ocl']) → str | None

kernelfoundry.algorithm.utils.extract_code.extract_last_code(output_string: str, code_language_types: list[str]) → str | None

Extract last code block from model output, specified by code_language_type

kernelfoundry.algorithm.utils.extract_code.extract_code_blocks(text, code_language_types: list[str]) → str

Extract all code blocks from text, combine them to return as a single string

kernelfoundry.algorithm.utils.extract_code.apply_diff(original_code: str, diff_text: str) → str

Apply a diff to the original code

Parameters:

• original_code – Original source code

• diff_text – Diff in the SEARCH/REPLACE format

Returns:

Modified code

kernelfoundry.algorithm.utils.extract_code.extract_diffs(diff_text: str) → List[Tuple[str, str]]

Extract diff blocks from the diff text

Parameters:

diff_text – Diff in the SEARCH/REPLACE format

Returns:

List of tuples (search_text, replace_text)

kernelfoundry.algorithm.utils.extract_code.format_diff_summary(diff_blocks: List[Tuple[str, str]]) → str

Create a human-readable summary of the diff

Parameters:

diff_blocks – List of (search_text, replace_text) tuples

Returns:

Summary string

kernelfoundry.algorithm.utils.extract_code.replace_function_calls(text)

For templated kernels, add template_args to pytorch functional