Fix inference: better tokenizer init, gcode cleaning, and debug output

app.py

@@ -304,21 +304,29 @@ def get_model():
         try:
             # 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)
-            # 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
+            gcode_tokenizer = PreTrainedTokenizerFast(
+                tokenizer_file=tokenizer_path,
+                pad_token="<pad>",
+                unk_token="<unk>",
+                bos_token="<s>",
+                eos_token="</s>",
+            )
+            # Verify special tokens
             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}")
+            print(f"  BOS='{gcode_tokenizer.bos_token}' (id={gcode_tokenizer.bos_token_id})")
+            print(f"  EOS='{gcode_tokenizer.eos_token}' (id={gcode_tokenizer.eos_token_id})")
+            print(f"  PAD='{gcode_tokenizer.pad_token}' (id={gcode_tokenizer.pad_token_id})")
+            
+            # Test encode/decode
+            test = "G0 X100 Y200\nG1 X150 Y250"
+            enc = gcode_tokenizer.encode(test)
+            dec = gcode_tokenizer.decode(enc)
+            print(f"  Test encode: {len(enc)} tokens")
+            print(f"  Test decode: '{dec[:50]}...'")
         except Exception as e:
             print(f"Failed to load custom tokenizer: {e}")
+            import traceback
+            traceback.print_exc()
             # Fallback to T5 tokenizer
             gcode_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
             print("Using fallback T5 tokenizer")
@@ -341,6 +349,72 @@ def get_model():
 # GCODE PROCESSING
 # ============================================================================
 
+def clean_gcode(gcode: str) -> str:
+    """Clean up generated gcode - fix formatting, remove garbage."""
+    
+    # Replace <newline> tokens with actual newlines
+    gcode = gcode.replace("<newline>", "\n")
+    
+    # If no/few newlines, split on command boundaries
+    if gcode.count("\n") < 10:
+        # Split before each gcode command
+        gcode = re.sub(r'([GM]\d+)', r'\n\1', gcode)
+    
+    # Clean up each line
+    cleaned_lines = []
+    seen_coords = set()  # Track to detect stuck coordinates
+    
+    for line in gcode.split("\n"):
+        line = line.strip()
+        if not line:
+            continue
+            
+        # Skip garbage/metadata lines
+        if line.lower() in ["dcode", "gcode", "code", "output"]:
+            continue
+        if line.startswith("Source:") or line.startswith(";Generated"):
+            continue
+        if line.startswith("Workarea:") or line.startswith("Algorithm:"):
+            continue
+            
+        # Fix malformed coordinates like X-X-X-100 or X-361.X-390
+        line = re.sub(r'X-X-X-', 'X-', line)
+        line = re.sub(r'X-X-', 'X-', line)
+        line = re.sub(r'X-\d+\.X-', 'X-', line)
+        line = re.sub(r'Y-Y-Y-', 'Y-', line)
+        line = re.sub(r'Y-Y-', 'Y-', line)
+        line = re.sub(r'Y-\d+\.Y-', 'Y-', line)
+        
+        # Fix missing spaces: G1X -> G1 X
+        line = re.sub(r'(G[01])X', r'\1 X', line)
+        line = re.sub(r'(G[01])Y', r'\1 Y', line)
+        
+        # Extract coordinates to check for stuck positions
+        x_match = re.search(r'X([-\d.]+)', line)
+        y_match = re.search(r'Y([-\d.]+)', line)
+        if x_match and y_match:
+            try:
+                coord = (round(float(x_match.group(1)), 1), round(float(y_match.group(1)), 1))
+                if coord in seen_coords:
+                    # Skip if we've seen this exact coordinate recently
+                    if len(seen_coords) > 5:
+                        continue
+                seen_coords.add(coord)
+                # Keep only last 50 coords
+                if len(seen_coords) > 50:
+                    seen_coords = set(list(seen_coords)[-50:])
+            except ValueError:
+                pass
+        
+        # Only keep lines starting with valid gcode commands
+        if line and line[0] in "GMgm;":
+            cleaned_lines.append(line)
+    
+    result = "\n".join(cleaned_lines)
+    print(f"Cleaned gcode: {len(cleaned_lines)} lines")
+    return result
+
+
 def validate_gcode(gcode: str) -> str:
     """Clamp coordinates to machine bounds."""
     lines = []
@@ -539,41 +613,53 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
         with torch.no_grad():
             batch_size = latent.shape[0]
             
-            # Start token - use BOS for v3, semicolon for v2
+            # Get proper token IDs
+            bos_id = gcode_tokenizer.bos_token_id
+            eos_id = gcode_tokenizer.eos_token_id  
+            pad_id = gcode_tokenizer.pad_token_id
+            
+            # For v3, start with BOS token; for v2, encode gcode header
             if is_v3:
-                start_id = gcode_tokenizer.bos_token_id if gcode_tokenizer.bos_token_id is not None else 1
+                # Use the gcode header as the starting prompt
+                start_text = "G21\nG90\nM280 P0 S90\nG28\n"
+                start_tokens = gcode_tokenizer.encode(start_text, add_special_tokens=False)
+                if bos_id is not None:
+                    start_tokens = [bos_id] + start_tokens
+                input_ids = torch.tensor([start_tokens], dtype=torch.long, device=device)
             else:
                 start_tokens = gcode_tokenizer.encode(";", add_special_tokens=False)
                 start_id = start_tokens[0] if start_tokens else 0
+                input_ids = torch.tensor([[start_id]], dtype=torch.long, device=device)
             
-            print(f"Starting generation with token id: {start_id}")
-            input_ids = torch.tensor([[start_id]], dtype=torch.long, device=device)
+            print(f"Starting with {input_ids.shape[1]} tokens, BOS={bos_id}, EOS={eos_id}")
             
-            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
+            max_gen = min(max_tokens, gcode_decoder.config.max_seq_len - input_ids.shape[1])
             
-            # Track generated content for repetition detection
+            # Track for repetition detection
             recent_tokens = []
-            recent_coords = []
-            repetition_window = 30
             
             for step in range(max_gen):
                 logits = gcode_decoder(latent, input_ids)
                 next_logits = logits[:, -1, :] / temperature
                 
-                # Repetition penalty - reduce probability of recent tokens
+                # Suppress pad and unk tokens
+                if pad_id is not None:
+                    next_logits[:, pad_id] = float('-inf')
+                next_logits[:, 1] = float('-inf')  # <unk>
+                
+                # Repetition penalty
                 if recent_tokens:
-                    for token_id in set(recent_tokens[-repetition_window:]):
-                        next_logits[:, token_id] *= 0.6  # Stronger penalty
+                    for token_id in set(recent_tokens[-30:]):
+                        next_logits[:, token_id] *= 0.7
                 
-                # Top-k + Top-p sampling for better coherence
-                top_k = 40
-                top_p = 0.9
+                # Top-k + Top-p sampling
+                top_k = 50
+                top_p = 0.92
                 
                 # Top-k filtering
                 top_k_logits, top_k_indices = torch.topk(next_logits, top_k, dim=-1)
                 
-                # Top-p filtering within top-k
+                # Top-p filtering
                 sorted_logits, sorted_idx = torch.sort(top_k_logits, descending=True, dim=-1)
                 cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
                 sorted_indices_to_remove = cumulative_probs > top_p
@@ -584,66 +670,34 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
                 probs = torch.softmax(sorted_logits, dim=-1)
                 sampled_idx = torch.multinomial(probs, num_samples=1)
                 
-                # Map back to vocabulary indices
                 next_token = top_k_indices.gather(-1, sorted_idx.gather(-1, sampled_idx))
                 input_ids = torch.cat([input_ids, next_token], dim=1)
                 recent_tokens.append(next_token.item())
                 
+                # Debug first few tokens
+                if step < 5:
+                    tok_str = gcode_tokenizer.decode([next_token.item()])
+                    print(f"  Step {step}: token={next_token.item()}, str='{tok_str}'")
+                
                 # Check EOS
-                if next_token.item() == eos_id:
+                if eos_id is not None and 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)) < 4:  # Less than 4 unique tokens in last 20
-                        print(f"Stopping due to token repetition at step {step}")
+                # Early stop on repetition
+                if len(recent_tokens) > 30:
+                    if len(set(recent_tokens[-30:])) < 5:
+                        print(f"Stopping due to repetition at step {step}")
                         break
             
-            print(f"Generated {input_ids.shape[1]} tokens")
+            print(f"Generated {input_ids.shape[1]} total tokens")
             
-            # Decode - skip special tokens
+            # Decode
             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 from <newline> token
-            if is_v3:
-                gcode = gcode.replace("<newline>", "\n")
-            
-            # 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"):
-                line = line.strip()
-                # Skip empty lines and metadata
-                if not line:
-                    continue
-                if line.startswith("Source:") or line.startswith(";Generated"):
-                    continue
-                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"Raw decoded (first 300 chars): {repr(gcode[:300])}")
             
-            gcode = "\n".join(filtered_lines)
-            print(f"Filtered gcode: {len(filtered_lines)} lines, {len(gcode)} chars")
+            # Clean up the gcode
+            gcode = clean_gcode(gcode)
         
         gcode = validate_gcode(gcode)
         line_count = len([l for l in gcode.split("\n") if l.strip()])