import json

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_str = ''.join(source)
            if 'target_modules="all-linear"' in source_str:
                source_str = source_str.replace('target_modules="all-linear",', 'target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],')
                # Split back while preserving newlines
                cell['source'] = [s + '\n' for s in source_str.split('\n')]
                if cell['source'] and cell['source'][-1].endswith('\n') and not source_str.endswith('\n'):
                    cell['source'][-1] = cell['source'][-1][:-1]
                modifications += 1
        elif isinstance(source, str):
            if 'target_modules="all-linear"' in source:
                cell['source'] = source.replace('target_modules="all-linear",', 'target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],')
                modifications += 1

if modifications > 0:
    with open(path, 'w', encoding='utf-8') as f:
        json.dump(nb, f)
    print(f'Successfully applied {modifications} fixes to target_modules.')
else:
    print('No changes needed or could not find targets.')