一个将代码块统一添加代码行号的脚本
在 LLM Agent 训练中，有时存在需要通过代码行号进行补全的方法。
这个脚本给任意给定的代码行统一添加代码行号。
import json
import re
import argparse
def add_line_numbers_to_input(input_jsonl_path, output_jsonl_path):
with open(input_jsonl_path, 'r', encoding='utf-8') as fin, \
open(output_jsonl_path, 'w', encoding='utf-8') as fout:
for line in fin:
try:
data = json.loads(line)
if "input" in data:
data["input"] = process_code_block(data["input"])
fout.write(json.dumps(data, ensure_ascii=False) + '\n')
except json.JSONDecodeError:
print(f"JSON 解码错误: {line}")
def process_code_block(code_str):
# 匹配代码块（包含```c和```， 如果需要匹配其他语言可以之后更改）
code_block_match = re.search(r'(```c\n)(.*?)(\n```)', code_str, re.DOTALL)
if not code_block_match:
print("未找到符合条件的代码块")
return code_str
# 获取代码块各部分
start_marker = code_block_match.group(1) # ```c\n
code_content = code_block_match.group(2) # 实际代码内容
end_marker = code_block_match.group(3) # \n```
# 处理行号（保留原始缩进）
code_lines = code_content.split('\n')
number_lines = []
line_num = 1
for code_line in code_lines:
number_lines.append(f"{line_num}: {code_line}")
line_num += 1
# 重新建立代码块
processed_code = (
code_str[:code_block_match.start()] + # 保留代码块前的所有内容
start_marker + # ```c\n
'\n'.join(number_lines) + # 带行号的代码
end_marker + # \n```
code_str[code_block_match.end():] # 保留代码块后的所有内容
)
return processed_code
def main():
parser = argparse.ArgumentParser(description="Add line numbers to code blocks in a JSONL file.")
parser.add_argument('-i', '--input', required=True, help='Input JSONL file path')
parser.add_argument('-o', '--output', required=False, default='output.jsonl', help='Output JSONL filepath, default=output.jsonl')
args = parser.parse_args()
add_line_numbers_to_input(args.input, args.output)
print(f"Processed file saved to {args.output}")
if __name__ == '__main__':
main()
