#!/usr/bin/python3
"""Fix the broken docstrings in deepseek_v4.py"""

filepath = "/root/nvidia-meeting/deepseek-v4-quant/patches/deepseek_v4.py"

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

# Replace the problematic triple-quoted docstrings with comments
content = content.replace(
    '    def _convert_nvfp4_attention_to_fp8(self):\n        """Convert NVFP4 attention weights to FP8 format."""',
    '    def _convert_nvfp4_attention_to_fp8(self):  # Convert NVFP4 attention weights to FP8'
)

content = content.replace(
    '    def _convert_nvfp4_module_to_fp8(self, mod, e2m1_lut, fp8_max):\n        """Convert a single NVFP4 Linear module to FP8 format."""',
    '    def _convert_nvfp4_module_to_fp8(self, mod, e2m1_lut, fp8_max):  # Convert single NVFP4 module to FP8'
)

# Also remove any stray docstring lines
lines = content.split('\n')
new_lines = []
skip_next = False
for i, line in enumerate(lines):
    stripped = line.strip()
    if stripped == '"""' or stripped == "'''":
        continue  # Skip standalone triple-quote lines
    if '"""Convert NVFP4' in line or '"""Convert a single' in line:
        continue  # Skip the one-line docstrings that might remain
    new_lines.append(line)

with open(filepath, 'w') as f:
    f.write('\n'.join(new_lines))

# Verify syntax
import ast
try:
    ast.parse('\n'.join(new_lines))
    print("Syntax OK")
except SyntaxError as e:
    print(f"Syntax error: {e}")