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.')