# Patch for flamingo_lm.py line 177
# Change: return super().forward(**kwargs)
# To: return super().forward(**{k: v for k, v in kwargs.items() if k != 'labels'})

import sys
original_file = '/home/user/app/open_flamingo/src/flamingo_lm.py'

with open(original_file, 'r') as f:
    content = f.read()

# Replace the problematic line
old_line = "return super().forward(**kwargs)  # Call the other parent's forward method"
new_line = "return super().forward(**{k: v for k, v in kwargs.items() if k != 'labels'})  # Remove labels"

if old_line in content:
    content = content.replace(old_line, new_line)
    with open(original_file, 'w') as f:
        f.write(content)
    print("✅ Patched flamingo_lm.py")
else:
    print("⚠️  Line not found, trying alternative...")
    # Try without comment
    old_line = "return super().forward(**kwargs)"
    new_line = "return super().forward(**{k: v for k, v in kwargs.items() if k != 'labels'})"
    content = content.replace(old_line, new_line)
    with open(original_file, 'w') as f:
        f.write(content)
    print("✅ Patched flamingo_lm.py (alternative)")