Fix gcode tokenizer config and add post-processing to restore newlines

app.py

@@ -305,8 +305,20 @@ def get_model():
             # Try loading custom tokenizer from v3 model
             tokenizer_path = hf_hub_download("twarner/dcode-sd-gcode-v3", "gcode_tokenizer/tokenizer.json")
             gcode_tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_path)
-            print("Loaded custom gcode tokenizer")
-        except Exception:
+            # Ensure special tokens are set
+            if gcode_tokenizer.pad_token is None:
+                gcode_tokenizer.pad_token = "<pad>"
+                gcode_tokenizer.pad_token_id = 0
+            if gcode_tokenizer.bos_token is None:
+                gcode_tokenizer.bos_token = "<s>"
+                gcode_tokenizer.bos_token_id = 1
+            if gcode_tokenizer.eos_token is None:
+                gcode_tokenizer.eos_token = "</s>"
+                gcode_tokenizer.eos_token_id = 2
+            print(f"Loaded custom gcode tokenizer (vocab={gcode_tokenizer.vocab_size})")
+            print(f"  BOS={gcode_tokenizer.bos_token_id}, EOS={gcode_tokenizer.eos_token_id}, PAD={gcode_tokenizer.pad_token_id}")
+        except Exception as e:
+            print(f"Failed to load custom tokenizer: {e}")
             # Fallback to T5 tokenizer
             gcode_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
             print("Using fallback T5 tokenizer")
@@ -365,19 +377,21 @@ def gcode_to_svg(gcode: str) -> str:
     x, y = 0.0, 0.0
     pen_down = False
 
-    # Split on newlines, newline tokens, or command boundaries
-    lines = []
     # Replace newline tokens with actual newlines
     gcode = gcode.replace("<newline>", "\n")
     
-    for line in gcode.replace(";", "\n;").split("\n"):
-        line = line.strip()
-        if not line:
+    # Split concatenated gcode into separate commands
+    # First split on explicit newlines
+    lines = []
+    for raw_line in gcode.split("\n"):
+        raw_line = raw_line.strip()
+        if not raw_line:
             continue
-        parts = re.split(r'(?=[GM]\d)', line)
+        # Split on command boundaries (G0, G1, M280, etc)
+        parts = re.split(r'(?=[GM]\d)', raw_line)
         for part in parts:
             part = part.strip()
-            if part and not part.startswith(";"):
+            if part and not part.startswith(";") and part[0] in "GMgm":
                 lines.append(part)
     
     for line in lines:
@@ -525,22 +539,23 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
         with torch.no_grad():
             batch_size = latent.shape[0]
             
-            # Start token
+            # Start token - use BOS for v3, semicolon for v2
             if is_v3:
-                # V3 uses custom tokenizer with BOS
-                start_id = gcode_tokenizer.bos_token_id or 0
+                start_id = gcode_tokenizer.bos_token_id if gcode_tokenizer.bos_token_id is not None else 1
             else:
-                # V2 uses semicolon as start
                 start_tokens = gcode_tokenizer.encode(";", add_special_tokens=False)
-                start_id = start_tokens[0] if start_tokens else gcode_tokenizer.pad_token_id
+                start_id = start_tokens[0] if start_tokens else 0
             
+            print(f"Starting generation with token id: {start_id}")
             input_ids = torch.tensor([[start_id]], dtype=torch.long, device=device)
             
             max_gen = min(max_tokens, gcode_decoder.config.max_seq_len - 1)
+            eos_id = gcode_tokenizer.eos_token_id if gcode_tokenizer.eos_token_id is not None else 2
             
             # Track generated content for repetition detection
             recent_tokens = []
-            repetition_window = 50
+            recent_coords = []
+            repetition_window = 30
             
             for step in range(max_gen):
                 logits = gcode_decoder(latent, input_ids)
@@ -549,11 +564,11 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
                 # Repetition penalty - reduce probability of recent tokens
                 if recent_tokens:
                     for token_id in set(recent_tokens[-repetition_window:]):
-                        next_logits[:, token_id] *= 0.7
+                        next_logits[:, token_id] *= 0.6  # Stronger penalty
                 
                 # Top-k + Top-p sampling for better coherence
-                top_k = 50
-                top_p = 0.85
+                top_k = 40
+                top_p = 0.9
                 
                 # Top-k filtering
                 top_k_logits, top_k_indices = torch.topk(next_logits, top_k, dim=-1)
@@ -575,33 +590,60 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
                 recent_tokens.append(next_token.item())
                 
                 # Check EOS
-                if next_token.item() == gcode_tokenizer.eos_token_id:
+                if next_token.item() == eos_id:
+                    print(f"Hit EOS at step {step}")
                     break
                 
                 # Early stop on excessive repetition
                 if len(recent_tokens) > 20:
                     last_20 = recent_tokens[-20:]
-                    if len(set(last_20)) < 5:  # Less than 5 unique tokens in last 20
-                        print("Stopping due to repetition")
+                    if len(set(last_20)) < 4:  # Less than 4 unique tokens in last 20
+                        print(f"Stopping due to token repetition at step {step}")
                         break
             
             print(f"Generated {input_ids.shape[1]} tokens")
+            
+            # Decode - skip special tokens
             gcode = gcode_tokenizer.decode(input_ids[0], skip_special_tokens=True)
+            print(f"Raw decoded (first 200): {gcode[:200]}")
             
-            # Post-process for v3: restore newlines
+            # Post-process for v3: restore newlines from <newline> token
             if is_v3:
                 gcode = gcode.replace("<newline>", "\n")
             
-            # Filter out training metadata lines
+            # If still no newlines, try to split on command boundaries
+            if "\n" not in gcode or gcode.count("\n") < 5:
+                print("No newlines found, splitting on command boundaries...")
+                # Split before G0, G1, G28, M280 commands
+                gcode = re.sub(r'(G0\s)', r'\n\1', gcode)
+                gcode = re.sub(r'(G1\s)', r'\n\1', gcode)
+                gcode = re.sub(r'(G1X)', r'\nG1 X', gcode)
+                gcode = re.sub(r'(G0X)', r'\nG0 X', gcode)
+                gcode = re.sub(r'(G28)', r'\nG28', gcode)
+                gcode = re.sub(r'(G21)', r'\nG21', gcode)
+                gcode = re.sub(r'(G90)', r'\nG90', gcode)
+                gcode = re.sub(r'(M280)', r'\nM280', gcode)
+                # Split on F speed values that are followed by another command
+                gcode = re.sub(r'(F\d+)(G)', r'\1\n\2', gcode)
+                gcode = re.sub(r'(F\d+)(M)', r'\1\n\2', gcode)
+            
+            # Filter out training metadata and garbage lines
             filtered_lines = []
             for line in gcode.split("\n"):
-                # Skip metadata headers from training data
+                line = line.strip()
+                # Skip empty lines and metadata
+                if not line:
+                    continue
                 if line.startswith("Source:") or line.startswith(";Generated"):
                     continue
-                filtered_lines.append(line)
-            gcode = "\n".join(filtered_lines)
+                if line.lower() in ["dcode", "gcode", "code"]:  # Skip garbage words
+                    continue
+                # Only keep lines that look like gcode (start with G, M, or ;)
+                if line[0] in "GMgm;":
+                    filtered_lines.append(line)
             
-            print(f"Decoded gcode length: {len(gcode)} chars")
+            gcode = "\n".join(filtered_lines)
+            print(f"Filtered gcode: {len(filtered_lines)} lines, {len(gcode)} chars")
         
         gcode = validate_gcode(gcode)
         line_count = len([l for l in gcode.split("\n") if l.strip()])