File size: 1,943 Bytes
10ff0db | 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 | import json
import io
path = r'c:\Users\vaibhav patil\.gemini\antigravity\scratch\financial-doc-extractor\scripts\financial-doc-training (1).ipynb'
with open(path, 'r', encoding='utf-8') as f:
nb = json.load(f)
modifications = 0
for cell in nb.get('cells', []):
if cell['cell_type'] == 'code':
source = cell['source']
if isinstance(source, list):
source = ''.join(source)
if 'import os' in source and 'import torch' in source and 'CUDA_VISIBLE_DEVICES' not in source:
source = source.replace('import os\n', 'import os\nos.environ["CUDA_VISIBLE_DEVICES"] = "0"\n')
modifications += 1
if 'target_modules=["q_proj", "v_proj"],' in source:
source = source.replace('target_modules=["q_proj", "v_proj"],', 'target_modules="all-linear",')
source = source.replace('# HARDCODING to 2 layers! Cuts activation memory by 60%', '# FIXED: "all-linear" enables Unsloth memory efficient fused Triton kernels.')
modifications += 1
if 'TrainingArguments(' in source and 'fp16' not in source:
source = source.replace('report_to="none",', 'report_to="none",\n fp16=not torch.cuda.is_bf16_supported(),\n bf16=torch.cuda.is_bf16_supported(),')
modifications += 1
if isinstance(cell['source'], list):
# Split back safely
cell['source'] = [s + '\n' for s in source.split('\n')]
if cell['source'] and cell['source'][-1].endswith('\n') and not source.endswith('\n'):
cell['source'][-1] = cell['source'][-1][:-1]
else:
cell['source'] = source
if modifications > 0:
with open(path, 'w', encoding='utf-8') as f:
json.dump(nb, f)
print(f'Successfully applied {modifications} fixes to the notebook.')
else:
print('No changes needed or could not find targets.')
|