from pathlib import Path
import ast

SCRIPT = Path("trainNLPRAGV2.py")

if not SCRIPT.exists():
    raise FileNotFoundError("trainNLPRAGV2.py introuvable dans le dossier courant.")

text = SCRIPT.read_text(encoding="utf-8")

tree = ast.parse(text)
target = None

for node in tree.body:
    if isinstance(node, ast.Assign):
        for t in node.targets:
            if isinstance(t, ast.Name) and t.id == "DATASET_SPECS":
                target = node
                break
    if target:
        break

if target is None:
    raise RuntimeError("Impossible de trouver DATASET_SPECS dans trainNLPRAGV2.py")

lines = text.splitlines()

new_block = '''DATASET_SPECS = [
    {
        "label": "boolq",
        "name": "google/boolq",
        "config": None,
        "split": "train",
        "kind": "boolq",
        "text_field": "passage",
        "group": "single",
        "trust_remote_code": False,
    }
]'''

start = target.lineno - 1
end = target.end_lineno

new_lines = lines[:start] + new_block.splitlines() + lines[end:]
SCRIPT.write_text("\\n".join(new_lines) + "\\n", encoding="utf-8")

print("OK : DATASET_SPECS remplacé par google/boolq uniquement.")