Update app.py

app.py

@@ -353,31 +353,80 @@ def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p
         
         generation_time = time.time() - start_time
         
-        # Decode all sequences
+        # Decode all sequences with error handling
         generated_codes = []
-        for output in outputs:
-            generated = loaded_tokenizer.decode(output, skip_special_tokens=False)
-            
-            # Extract code part
-            if '<CODE>' in generated:
-                code = generated.split('<CODE>')[-1].strip()
-                # Remove special tokens
-                code = code.replace('<PAD>', '').replace('<SEP>', '').strip()
-            else:
-                code = generated
-            
-            generated_codes.append(code)
+        for i, output in enumerate(outputs):
+            try:
+                # Ensure output is valid tensor and contains valid token IDs
+                if output is None:
+                    continue
+                
+                # Convert to list and filter out None values
+                if hasattr(output, 'tolist'):
+                    token_ids = output.tolist()
+                else:
+                    token_ids = output
+                
+                # Filter out None values and ensure all are integers
+                valid_tokens = []
+                for token in token_ids:
+                    if token is not None and isinstance(token, (int, float)):
+                        valid_tokens.append(int(token))
+                
+                if not valid_tokens:
+                    generated_codes.append(f"# Generation {i+1} failed: No valid tokens")
+                    continue
+                
+                # Decode with skip_special_tokens=True for cleaner output
+                try:
+                    generated = loaded_tokenizer.decode(valid_tokens, skip_special_tokens=False)
+                except Exception as decode_error:
+                    # Fallback: try with skip_special_tokens=True
+                    try:
+                        generated = loaded_tokenizer.decode(valid_tokens, skip_special_tokens=True)
+                    except Exception as decode_error2:
+                        # Last resort: convert tokens to string manually
+                        generated = f"# Decode failed: {str(decode_error2)}"
+                
+                # Handle None result from decode
+                if generated is None:
+                    generated = f"# Generation {i+1}: Decode returned None"
+                
+                # Extract code part
+                if '<CODE>' in generated:
+                    code = generated.split('<CODE>')[-1].strip()
+                    # Remove special tokens
+                    code = code.replace('<PAD>', '').replace('<SEP>', '').replace('</s>', '').replace('<s>', '').strip()
+                else:
+                    code = generated.strip()
+                
+                # Ensure we have some content
+                if not code or code.isspace():
+                    code = f"# Generated sequence {i+1} was empty"
+                
+                generated_codes.append(code)
+                
+            except Exception as decode_error:
+                # Handle any other decoding errors
+                error_msg = f"# Error decoding sequence {i+1}: {str(decode_error)}"
+                generated_codes.append(error_msg)
+        
+        # Ensure we have at least one result
+        if not generated_codes:
+            generated_codes = ["# No valid generations produced"]
         
         # Use the first generated code as primary output
-        primary_code = generated_codes[0]
+        primary_code = generated_codes[0] if generated_codes else "# No code generated"
         
         # Calculate metrics if reference code is provided
         metrics_output = ""
         bleu_output = ""
         
-        if reference_code and reference_code.strip():
-            # Calculate BLEU scores
-            bleu_1, bleu_2, bleu_3, bleu_4 = calculate_bleu_score(reference_code, primary_code)
+        if reference_code and reference_code.strip() and not primary_code.startswith('#'):
+            # Only calculate metrics if we have valid generated code (not error messages)
+            try:
+                # Calculate BLEU scores
+                bleu_1, bleu_2, bleu_3, bleu_4 = calculate_bleu_score(reference_code, primary_code)
             
             bleu_output = f"""📊 BLEU Scores:
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
@@ -410,6 +459,13 @@ def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p
 📝 Sequences Generated: {num_sequences}
 🔢 Output Length: {len(primary_code)} characters
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+"""
+            except Exception as metrics_error:
+                metrics_output = f"""⚠️ Metrics calculation failed: {str(metrics_error)}
+
+⏱️ Generation Time: {generation_time:.2f}s
+📝 Sequences Generated: {num_sequences}
+🔢 Output Length: {len(primary_code)} characters
 """
         else:
             metrics_output = f"""⏱️ Generation Time: {generation_time:.2f}s