Update app.py

app.py

@@ -5,22 +5,29 @@ import torch
 from threading import Thread
 
 veri_model_path = "nyu-dice-lab/VeriThoughts-Reasoning-7B"
-
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-# Try loading the model with explicit error handling
+# Try loading the model with KV caching (no flash attention or quantization)
 try:
+    print("Loading tokenizer...")
     veri_tokenizer = AutoTokenizer.from_pretrained(veri_model_path)
-
+    
+    # Set pad token if not exists
+    if veri_tokenizer.pad_token is None:
+        veri_tokenizer.pad_token = veri_tokenizer.eos_token
+    
+    print("Loading model with KV caching...")
     veri_model = AutoModelForCausalLM.from_pretrained(
         veri_model_path, 
-        device_map="auto", 
-        torch_dtype="auto",
+        device_map="auto" if torch.cuda.is_available() else None,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
         trust_remote_code=True,
-        use_cache=True,  # Enable KV caching
- #       attn_implementation="flash_attention_2" if torch.cuda.is_available() else None
+        use_cache=True,  # Enable KV caching for faster generation
+        low_cpu_mem_usage=True
     )
-
+    
+    print("Model loaded successfully with KV caching!")
+    
 except Exception as e:
     print(f"Model loading error: {e}")
     veri_model = None
@@ -33,100 +40,260 @@ def truncate_at_code_end(text):
         end_index = text.find("CODE END") + len("CODE END")
         return text[:end_index].strip()
     return text.strip()
-    
+
 def generate_response(user_message, history):
+    """Non-streaming generation for quick responses"""
     if not veri_model or not veri_tokenizer:
         return history + [["Error", "Model not loaded properly"]]
     
     if not user_message.strip():
         return history
-        
-    # Simple generation without streaming first
-    system_message = "You are VeriThoughts, a helpful assistant that thinks step by step. You are finetuned from a Qwen model, created by Alibaba Cloud. If you are asked a Verilog question, make sure your input and output interface has the same names as described in the question. If you are asked to generate code, please start your Verilog code with CODE BEGIN and end with CODE END."
     
+    system_message = "You are VeriThoughts, a helpful assistant that thinks step by step to answer Verilog coding questions. Make sure your input and output interface has the same names as described in the question. Please start your Verilog code with CODE BEGIN and end with CODE END."
+    
+    # Create conversation history (limit to last 3 exchanges for memory efficiency)
     conversation = f"System: {system_message}\n"
     recent_history = history[-3:] if len(history) > 3 else history
-  
+    
     for h in recent_history:
         conversation += f"User: {h[0]}\nAssistant: {h[1]}\n"
     conversation += f"User: {user_message}\nAssistant:"
     
+    # Tokenize input
     inputs = veri_tokenizer(
-        conversation,                     
+        conversation, 
         return_tensors="pt", 
         truncation=True, 
-        max_length=8192,
- #       padding=True,
- #       return_attention_mask=True
+        max_length=4096,
+        padding=True
     ).to(device)
     
+    # Generate with KV caching
     with torch.no_grad():
         outputs = veri_model.generate(
             **inputs,
-            max_new_tokens=20000,
+            max_new_tokens=1024,
             temperature=0.6,
             top_p=0.95,
             do_sample=True,
-            use_cache=True,      # Enable KV caching for faster generation
-            repetition_penalty=1.1,  # Reduce repetition
- #           length_penalty=1.0,
- #           early_stopping=True,     # Stop early when appropriate
- #           num_beams=1,            # Greedy search for speed
- #           pad_token_id=veri_tokenizer.eos_token_id
+            pad_token_id=veri_tokenizer.pad_token_id,
+            eos_token_id=veri_tokenizer.eos_token_id,
+            use_cache=True,  # KV caching for speed
+            repetition_penalty=1.1,
+            early_stopping=True
         )
     
+    # Decode response
     response = veri_tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-
+    
     # Truncate at CODE END to remove repetitive content
-    # response = truncate_at_code_end(response)
+    response = truncate_at_code_end(response)
     
-
+    # Clean up GPU memory
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
     
-    # Return updated history
-    return history + [[user_message, response.strip()]]
+    return history + [[user_message, response]]
 
-# Create minimal interface
-with gr.Blocks(
-    title="VeriThoughts-7B Chatbot",
-    css="""
-    .gradio-container {
-        max-width: 1200px !important;
-    }
-    .chat-message {
-        font-size: 14px;
-    }
-    """
-) as demo:
-    gr.Markdown(
-        """
-        # 🤖 VeriThoughts-7B Chatbot
+@spaces.GPU(duration=120)
+def generate_response_streaming(user_message, history):
+    """Streaming generation for real-time response display"""
+    if not veri_model or not veri_tokenizer:
+        yield history + [["Error", "Model not loaded properly"]]
+        return
+    
+    if not user_message.strip():
+        yield history
+        return
+    
+    system_message = "You are VeriThoughts, a helpful assistant that thinks step by step. You are finetuned from a Qwen model, created by Alibaba Cloud, to answer Verilog coding questions. Make sure your input and output interface has the same names as described in the question. Please start your Verilog code with CODE BEGIN and end with CODE END."
+    
+    # Create conversation history (limit for memory efficiency)
+    conversation = f"System: {system_message}\n"
+    recent_history = history[-3:] if len(history) > 3 else history
+    
+    for h in recent_history:
+        conversation += f"User: {h[0]}\nAssistant: {h[1]}\n"
+    conversation += f"User: {user_message}\nAssistant:"
+    
+    try:
+        # Tokenize input
+        inputs = veri_tokenizer(
+            conversation, 
+            return_tensors="pt", 
+            truncation=True, 
+            max_length=2048,
+            padding=True
+        ).to(device)
+        
+        # Setup streaming
+        streamer = TextIteratorStreamer(
+            veri_tokenizer, 
+            skip_prompt=True, 
+            skip_special_tokens=True,
+            timeout=30.0
+        )
+        
+        # Generation parameters with KV caching
+        generation_kwargs = {
+            **inputs,
+            "max_new_tokens": 4096,
+            "temperature": 0.6,
+            "top_p": 0.95,
+            "do_sample": True,
+            "pad_token_id": veri_tokenizer.pad_token_id,
+            "eos_token_id": veri_tokenizer.eos_token_id,
+            "use_cache": True,  # KV caching for faster streaming
+            "repetition_penalty": 1.1,
+            "streamer": streamer,
+            "early_stopping": True
+        }
+        
+        # Start generation in a separate thread
+        thread = Thread(target=veri_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        # Stream the response token by token
+        generated_text = ""
+        new_history = history + [[user_message, ""]]
+        code_end_reached = False
         
-        An AI assistant specialized in Verilog coding and digital design.
+        for new_text in streamer:
+            # Stop streaming if we've already reached CODE END
+            if code_end_reached:
+                break
+                
+            generated_text += new_text
+            
+            # Check if CODE END appears in the generated text
+            if "CODE END" in generated_text:
+                # Truncate at CODE END and mark as complete
+                generated_text = truncate_at_code_end(generated_text)
+                code_end_reached = True
+            
+            new_history[-1][1] = generated_text
+            yield new_history
+            
+            # Break early if CODE END was reached
+            if code_end_reached:
+                break
         
-        **Tips for better results:**
-        - Mention input/output port names clearly
-        - Ask for step-by-step explanations
-        """
+        # Ensure the thread completes
+        thread.join()
+        
+        # Final cleanup in case CODE END wasn't reached during streaming
+        if not code_end_reached:
+            final_text = truncate_at_code_end(generated_text)
+            new_history[-1][1] = final_text
+            yield new_history
+        
+    except Exception as e:
+        print(f"Streaming error: {e}")
+        error_history = history + [[user_message, f"Streaming error: {str(e)}"]]
+        yield error_history
+    
+    finally:
+        # Clean up GPU memory after generation
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+def clear_chat():
+    """Clear chat and clean up memory"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    return []
+
+# Create interface with soft theme
+with gr.Blocks(title="VeriThoughts-7B Chatbot") as demo:
+    gr.Markdown("# VeriThoughts-7B Chatbot")
+    gr.Markdown("*Optimized with KV caching for faster generation*")
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(
+                value=[], 
+                label="Chat", 
+                height=600,
+                show_label=False,
+                container=True
+            )
+            
+            with gr.Row():
+                msg = gr.Textbox(
+                    label="Your message", 
+                    placeholder="Ask me about Verilog design, syntax, or implementation...",
+                    lines=2,
+                    max_lines=5,
+                    scale=4
+                )
+                send_btn = gr.Button("Send", variant="primary", scale=1)
+        
+        with gr.Column(scale=1):
+            with gr.Group():
+                stream_btn = gr.Button("📡 Send (Streaming)", variant="secondary", size="sm")
+                clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary", size="sm")
+            
+            gr.Markdown(
+                """
+                ### 💡 Usage Tips
+                
+                **Send**: Quick response (max 1K tokens)
+                **Streaming**: Real-time response (max 2K tokens)
+                
+                ### ⚡ Optimizations Active
+                - **KV Caching**: Faster token generation
+                - **Memory Management**: Auto cleanup
+                - **Context Limiting**: Recent history only
+                
+                ### 🎯 Best Practices
+                - Be specific about Verilog requirements
+                - Mention input/output port names
+                - Ask for step-by-step explanations
+                - Clear chat periodically
+                """
+            )
+    
+    # Event handlers for regular send
+    submit_event = msg.submit(
+        fn=generate_response,
+        inputs=[msg, chatbot],
+        outputs=chatbot,
+        show_progress=True
+    ).then(
+        lambda: "",
+        inputs=None,
+        outputs=msg
     )
     
-    chatbot = gr.Chatbot(value=[], label="Chat")
-    msg = gr.Textbox(label="Your message", placeholder="Ask me about Verilog design, syntax, or implementation...")
-    clear = gr.Button("Clear")
+    send_btn.click(
+        fn=generate_response,
+        inputs=[msg, chatbot],
+        outputs=chatbot,
+        show_progress=True
+    ).then(
+        lambda: "",
+        inputs=None,
+        outputs=msg
+    )
     
-    # Simple event handling
-    msg.submit(
-        fn=generate_response, 
-        inputs=[msg, chatbot], 
-        outputs=chatbot
+    # Event handler for streaming
+    stream_btn.click(
+        fn=generate_response_streaming,
+        inputs=[msg, chatbot],
+        outputs=chatbot,
+        show_progress=True
     ).then(
-        lambda: "", 
-        inputs=None, 
+        lambda: "",
+        inputs=None,
         outputs=msg
     )
     
-    clear.click(lambda: [], outputs=chatbot)
+    # Clear chat handler
+    clear_btn.click(
+        fn=clear_chat,
+        inputs=None,
+        outputs=chatbot
+    )
 
-# Launch without ssr_mode parameter which might cause issues
+# Launch the app
 demo.launch(share=True)