Update app.py

app.py

@@ -734,63 +734,110 @@ def generate_text_phoenix(
         start_time = time.time()
         generated_ids = []
         
+        model.eval()  # ✅ Set to eval mode
+        
         with torch.no_grad():
             for step in range(max_new_tokens):
-                # Forward pass (now with lm_head)
-                outputs = model(input_ids=input_ids)
-                
-                # Get logits from lm_head
-                logits = outputs.logits[:, -1, :]  # [B, vocab_size]
-                
-                # Temperature sampling
-                if temperature > 0:
-                    probs = F.softmax(logits / temperature, dim=-1)
-                    next_token = torch.multinomial(probs, num_samples=1)
-                else:
-                    next_token = logits.argmax(dim=-1, keepdim=True)
-                
-                # Append
-                generated_ids.append(next_token.item())
-                input_ids = torch.cat([input_ids, next_token], dim=1)
-                
-                # Stop at EOS
-                if next_token.item() == tokenizer.eos_token_id:
-                    print(f"   Stopped at EOS token")
+                try:
+                    # Forward pass (now with lm_head)
+                    outputs = model(input_ids=input_ids)
+                    
+                    # Get logits from lm_head
+                    logits = outputs.logits[:, -1, :]  # [B, vocab_size]
+                    
+                    # ✅ Clamp logits to prevent numerical issues
+                    logits = torch.clamp(logits, min=-100, max=100)
+                    
+                    # Temperature sampling
+                    if temperature > 0.01:
+                        logits = logits / temperature
+                        probs = F.softmax(logits, dim=-1)
+                        
+                        # ✅ Check for NaN/Inf
+                        if torch.isnan(probs).any() or torch.isinf(probs).any():
+                            print(f"   ⚠️ NaN/Inf detected at step {step}, using greedy")
+                            next_token = logits.argmax(dim=-1, keepdim=True)
+                        else:
+                            # ✅ Add small epsilon to avoid zero probabilities
+                            probs = probs + 1e-10
+                            probs = probs / probs.sum(dim=-1, keepdim=True)
+                            next_token = torch.multinomial(probs, num_samples=1)
+                    else:
+                        next_token = logits.argmax(dim=-1, keepdim=True)
+                    
+                    next_token_id = next_token.item()
+                    
+                    # ✅ Validate token range
+                    if next_token_id < 0 or next_token_id >= model.config.vocab_size:
+                        print(f"   ⚠️ Invalid token {next_token_id}, stopping")
+                        break
+                    
+                    # Append
+                    generated_ids.append(next_token_id)
+                    input_ids = torch.cat([input_ids, next_token], dim=1)
+                    
+                    # ✅ Limit max sequence length
+                    if input_ids.shape[1] > 2048:
+                        print(f"   ⚠️ Max sequence length reached, stopping")
+                        break
+                    
+                    # Stop at EOS
+                    if next_token_id == tokenizer.eos_token_id:
+                        print(f"   ✅ Stopped at EOS token")
+                        break
+                    
+                    # Progress
+                    if (step + 1) % 10 == 0:
+                        print(f"   Generated {step + 1}/{max_new_tokens} tokens...")
+                        
+                except RuntimeError as e:
+                    print(f"   ❌ Runtime error at step {step}: {e}")
+                    if "CUDA" in str(e):
+                        print(f"   Stopping generation due to CUDA error")
+                    break
+                except Exception as e:
+                    print(f"   ❌ Error at step {step}: {e}")
                     break
-                
-                # Progress
-                if (step + 1) % 10 == 0:
-                    print(f"   Generated {step + 1}/{max_new_tokens} tokens...")
         
         elapsed = time.time() - start_time
         
-        # 5. 디코드
-        generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
-        full_text = prompt + generated_text
-        
-        # 6. 결과
+        # 6. 디코드
+        if len(generated_ids) == 0:
+            generated_text = "[No tokens generated]"
+            full_text = prompt
+        else:
+            try:
+                generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+                full_text = prompt + " " + generated_text
+            except Exception as e:
+                generated_text = f"[Decode error: {e}]"
+                full_text = prompt
+        
+        # 7. 결과
         output_md = f"""
 ## 📝 Generated Text
 
-**Prompt**: {prompt}
-
----
-
-**Generated**:
+**Prompt**: 
+```
+{prompt}
+```
 
+**Generated** ({len(generated_ids)} tokens):
+```
 {generated_text}
-
----
+```
 
 **Full Text**:
-
+```
 {full_text}
+```
 """
         
+        initial_tokens = input_ids.shape[1] - len(generated_ids)
         stats_md = f"""
 ## 📊 Generation Statistics
 
-- **Input tokens**: {input_ids.shape[1] - len(generated_ids)}
+- **Input tokens**: {initial_tokens}
 - **Generated tokens**: {len(generated_ids)}
 - **Total tokens**: {input_ids.shape[1]}
 - **Time**: {elapsed:.2f}s