Update app.py

app.py

@@ -377,56 +377,103 @@ def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p
                     generated_codes.append(f"# Generation {i+1} failed: No valid tokens")
                     continue
                 
-                # Decode with comprehensive error handling
+                # Decode with GPT-2 compatible handling
                 try:
-                    # First attempt: decode with skip_special_tokens=False
-                    generated = loaded_tokenizer.decode(valid_tokens, skip_special_tokens=False)
+                    # First attempt: standard decode with proper cleanup
+                    generated = loaded_tokenizer.decode(valid_tokens, skip_special_tokens=False, clean_up_tokenization_spaces=True)
                     
-                    # Check if decode returned None or contains None
+                    # GPT-2 specific: handle byte-level tokens properly
                     if generated is None:
                         raise ValueError("Tokenizer decode returned None")
                     
-                    # Check for None in the string (shouldn't happen but be safe)
-                    if 'None' in str(generated) or '\x00' in str(generated):
-                        raise ValueError("Decoded string contains invalid characters")
+                    # Clean up common GPT-2 artifacts
+                    generated = generated.replace('Ġ', ' ').replace('▁', ' ')  # Handle different space tokens
+                    generated = ' '.join(generated.split())  # Normalize whitespace
+                    
+                    # Check for gibberish (too many special characters)
+                    special_ratio = sum(1 for c in generated if not c.isalnum() and c not in ' \n\t.,;()[]{}+-*/=<>!&|^~%#@?:\'\"') / max(len(generated), 1)
+                    if special_ratio > 0.5:  # More than 50% special chars = likely gibberish
+                        raise ValueError("Decoded output appears to be gibberish")
                         
                 except Exception as decode_error:
                     # Second attempt: decode with skip_special_tokens=True
                     try:
-                        generated = loaded_tokenizer.decode(valid_tokens, skip_special_tokens=True)
+                        generated = loaded_tokenizer.decode(valid_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=True)
                         if generated is None:
                             raise ValueError("Tokenizer decode (skip_special) returned None")
-                        if 'None' in str(generated) or '\x00' in str(generated):
-                            raise ValueError("Decoded string contains invalid characters")
+                        
+                        # Clean up GPT-2 artifacts
+                        generated = generated.replace('Ġ', ' ').replace('▁', ' ')
+                        generated = ' '.join(generated.split())
+                        
+                        # Check for gibberish again
+                        special_ratio = sum(1 for c in generated if not c.isalnum() and c not in ' \n\t.,;()[]{}+-*/=<>!&|^~%#@?:\'\"') / max(len(generated), 1)
+                        if special_ratio > 0.5:
+                            raise ValueError("Decoded output still appears to be gibberish")
+                            
                     except Exception as decode_error2:
-                        # Third attempt: manual token-to-string conversion
+                        # Third attempt: manual byte-level decoding for GPT-2
                         try:
-                            # Convert tokens to string manually using vocab
-                            if hasattr(loaded_tokenizer, 'get_vocab'):
-                                vocab = loaded_tokenizer.get_vocab()
-                                inv_vocab = {v: k for k, v in vocab.items()}
-                                
-                                # Convert tokens to strings, skip unknown tokens
-                                token_strings = []
+                            # GPT-2 uses byte-level BPE, so we need to decode bytes properly
+                            if hasattr(loaded_tokenizer, 'byte_decoder'):
+                                # Use the tokenizer's byte decoder
+                                byte_tokens = []
                                 for token_id in valid_tokens:
-                                    if token_id in inv_vocab:
-                                        token_str = inv_vocab[token_id]
-                                        # Skip special tokens that might cause issues
-                                        if token_str not in ['<pad>', '<unk>', '<mask>', '<s>', '</s>', '<PAD>', '<SEP>', '<CODE>', '<PSEUDO>']:
-                                            token_strings.append(token_str)
-                                
-                                generated = ''.join(token_strings)
+                                    if token_id in loaded_tokenizer.decoder:
+                                        token_bytes = loaded_tokenizer.decoder[token_id]
+                                        if isinstance(token_bytes, bytes):
+                                            byte_tokens.append(token_bytes)
+                                        elif isinstance(token_bytes, str):
+                                            byte_tokens.append(token_bytes.encode('utf-8', errors='ignore'))
                                 
-                                if not generated or generated.isspace():
-                                    raise ValueError("Manual conversion produced empty string")
+                                if byte_tokens:
+                                    # Decode the byte sequence
+                                    full_bytes = b''.join(byte_tokens)
+                                    generated = full_bytes.decode('utf-8', errors='replace')
+                                    
+                                    # Clean up
+                                    generated = generated.replace('Ġ', ' ').replace('▁', ' ')
+                                    generated = ' '.join(generated.split())
+                                    
+                                    if not generated or generated.isspace():
+                                        raise ValueError("Byte decoding produced empty result")
+                                else:
+                                    raise ValueError("No valid byte tokens found")
                             else:
-                                raise ValueError("Tokenizer has no get_vocab method")
+                                raise ValueError("Tokenizer has no byte_decoder")
                                 
-                        except Exception as manual_error:
-                            # Final fallback: create a safe representation
-                            generated = f"# Decode failed: {str(decode_error)}\n# Manual conversion failed: {str(manual_error)}\n# Raw tokens: {valid_tokens[:10]}..."
+                        except Exception as byte_error:
+                            # Fourth attempt: fallback to vocab-based conversion
+                            try:
+                                if hasattr(loaded_tokenizer, 'get_vocab'):
+                                    vocab = loaded_tokenizer.get_vocab()
+                                    
+                                    # Convert tokens, handling byte-level tokens
+                                    text_parts = []
+                                    for token_id in valid_tokens:
+                                        if token_id in vocab:
+                                            token_text = vocab[token_id]
+                                            # Handle byte-level tokens (start with Ġ or ▁)
+                                            if token_text.startswith('Ġ'):
+                                                text_parts.append(' ' + token_text[1:])
+                                            elif token_text.startswith('▁'):
+                                                text_parts.append(' ' + token_text[1:])
+                                            else:
+                                                text_parts.append(token_text)
+                                    
+                                    generated = ''.join(text_parts)
+                                    generated = ' '.join(generated.split())  # Clean whitespace
+                                    
+                                    if not generated or generated.isspace():
+                                        raise ValueError("Vocab conversion produced empty result")
+                                else:
+                                    raise ValueError("Tokenizer has no get_vocab method")
+                                    
+                            except Exception as vocab_error:
+                                # Final fallback: show what we have
+                                generated = f"# Decode failed: {str(decode_error)}\n# Byte decode failed: {str(byte_error)}\n# Vocab decode failed: {str(vocab_error)}\n# Raw tokens: {valid_tokens[:10]}..."
                 
-                # Final safety check: ensure we have a string
+                # Final safety check
                 if not isinstance(generated, str):
                     generated = str(generated) if generated is not None else "# Decode returned non-string object"