Fix: keep newline tokens during decode, then convert to actual newlines

app.py

@@ -352,7 +352,7 @@ def get_model():
 def clean_gcode(gcode: str) -> str:
     """Clean up generated gcode - fix formatting, remove garbage."""
     
-    # Replace <newline> tokens with actual newlines
+    # Replace any remaining <newline> tokens
     gcode = gcode.replace("<newline>", "\n")
     
     # If no/few newlines, split on command boundaries
@@ -360,6 +360,10 @@ def clean_gcode(gcode: str) -> str:
         # Split before each gcode command
         gcode = re.sub(r'([GM]\d+)', r'\n\1', gcode)
     
+    # Add spaces after G0/G1 if missing: G0X -> G0 X
+    gcode = re.sub(r'(G[01])([XYZ])', r'\1 \2', gcode)
+    gcode = re.sub(r'(G[01])F', r'\1 F', gcode)
+    
     # Clean up each line
     cleaned_lines = []
     seen_coords = set()  # Track to detect stuck coordinates
@@ -692,8 +696,15 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
             
             print(f"Generated {input_ids.shape[1]} total tokens")
             
-            # Decode
-            gcode = gcode_tokenizer.decode(input_ids[0], skip_special_tokens=True)
+            # Decode WITHOUT skipping special tokens (so we keep <newline>)
+            gcode = gcode_tokenizer.decode(input_ids[0], skip_special_tokens=False)
+            
+            # Manually remove the special tokens we don't want, but keep <newline>
+            gcode = gcode.replace("<pad>", "").replace("<s>", "").replace("</s>", "").replace("<unk>", "")
+            
+            # Now convert <newline> to actual newlines
+            gcode = gcode.replace("<newline>", "\n")
+            
             print(f"Raw decoded (first 300 chars): {repr(gcode[:300])}")
             
             # Clean up the gcode