Fix code block extraction and increase completion length

app.py

@@ -29,6 +29,7 @@ check_import("transformers", lambda: __import__("transformers").__version__)
 check_import("datasets", lambda: __import__("datasets").__version__)
 check_import("peft", lambda: __import__("peft").__version__)
 check_import("trl", lambda: __import__("trl").__version__)
+check_import("huggingface_hub", lambda: __import__("huggingface_hub").__version__)
 
 try:
     from trl import GRPOConfig, GRPOTrainer
@@ -103,6 +104,7 @@ def get_status():
 
 
 def extract_code_block(text: str) -> str:
+    # Prefer closed fences
     pattern = r"```python\s*(.*?)```"
     matches = re.findall(pattern, text, re.DOTALL)
     if matches:
@@ -111,6 +113,18 @@ def extract_code_block(text: str) -> str:
     matches = re.findall(pattern, text, re.DOTALL)
     if matches:
         return matches[-1].strip()
+
+    # Handle unclosed fences (common when generation truncates)
+    if "```python" in text:
+        after = text.split("```python", 1)[1]
+        if "```" in after:
+            after = after.split("```", 1)[0]
+        return after.strip()
+    if "```" in text:
+        after = text.split("```", 1)[1]
+        if "```" in after:
+            after = after.split("```", 1)[0]
+        return after.strip()
     return text.strip()
 
 
@@ -591,7 +605,7 @@ def run_training(model_name, chunk_steps, max_total_steps, max_minutes, auto_con
                 save_steps=999999,
                 report_to="none",
                 remove_unused_columns=False,
-                max_completion_length=512,
+                max_completion_length=2048,
                 num_generations=4,
             )
 
@@ -627,7 +641,7 @@ def run_training(model_name, chunk_steps, max_total_steps, max_minutes, auto_con
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=256,
+                max_new_tokens=1024,
                 do_sample=True,
                 temperature=0.7,
                 top_p=0.9,