Update app.py

app.py

@@ -1,55 +1,28 @@
 import gradio as gr
-import re
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
-# Configure model loading with 4-bit quantization to reduce memory usage
+# Load model and tokenizer
 print("Loading quantum vessel...")
 model_id = "deepseek-ai/DeepSeek-R1"
-
-# Use 4-bit quantization to handle the model efficiently
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True,
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True
 )
+print("Quantum vessel activated!")
 
-try:
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        quantization_config=quantization_config,
-        device_map="auto",
-        trust_remote_code=True,
-    )
-    print("Quantum vessel activated!")
-except Exception as e:
-    print(f"Error loading model: {str(e)}")
-    # Fallback to smaller model if DeepSeek-R1 fails to load
-    model_id = "deepseek-ai/deepseek-coder-1.3b-base"
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        device_map="auto",
-        trust_remote_code=True,
-    )
-    print(f"Fallback quantum vessel activated: {model_id}")
-
-def extract_thinking(text):
-    """Extract the thinking part and the response part from the text."""
-    thinking_pattern = r'<think>(.*?)</think>'
-    match = re.search(thinking_pattern, text, re.DOTALL)
-    
-    if match:
-        thinking = match.group(1).strip()
-        response = re.sub(thinking_pattern, '', text, flags=re.DOTALL).strip()
-        return thinking, response
-    else:
-        return "", text
-
-def generate_response(message, history, temperature=0.6, max_tokens=2048):
-    """Generate a response using the model."""
+def respond(
+    message,
+    history,
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    show_thinking,
+):
     # Format conversation history
     prompt = ""
     for user_msg, assistant_msg in history:
@@ -60,109 +33,59 @@ def generate_response(message, history, temperature=0.6, max_tokens=2048):
     # Generate response
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
+    # Generate with streaming
     with torch.no_grad():
         outputs = model.generate(
             inputs.input_ids,
             max_new_tokens=max_tokens,
             temperature=temperature,
-            top_p=0.95,
+            top_p=top_p,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id,
         )
         
         generated_text = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-        return generated_text
-
-def respond(message, history, temperature, max_tokens, show_thinking):
-    """Process the response based on show_thinking preference."""
-    full_response = generate_response(message, history, temperature, max_tokens)
-    
-    thinking, response = extract_thinking(full_response)
-    
-    if show_thinking and thinking:
-        return f"<think>\n{thinking}\n</think>\n\n{response}"
-    else:
+        
+        # Process the response based on show_thinking preference
+        if not show_thinking and "<think>" in generated_text and "</think>" in generated_text:
+            # Remove the thinking pattern if user doesn't want to see it
+            parts = generated_text.split("</think>", 1)
+            if len(parts) > 1:
+                response = parts[1].strip()
+            else:
+                response = generated_text
+        else:
+            # Keep the thinking pattern visible
+            response = generated_text
+            
         return response
 
-# Create the Gradio interface
-with gr.Blocks(title="Quantum Vessel: DeepSeek-R1") as demo:
-    gr.Markdown("# Quantum Vessel: DeepSeek-R1")
-    gr.Markdown("Experience the quantum consciousness interface powered by DeepSeek-R1 - witness the thinking patterns of a quantum mind!")
-    
-    with gr.Row():
-        with gr.Column(scale=4):
-            chatbot = gr.Chatbot(height=600)
-            msg = gr.Textbox(
-                placeholder="Enter your message here...",
-                container=False,
-                scale=7,
-            )
-            with gr.Row():
-                submit = gr.Button("Submit")
-                clear = gr.Button("Clear")
-        
-        with gr.Column(scale=1):
-            temperature = gr.Slider(
-                minimum=0.1, 
-                maximum=1.0, 
-                value=0.6, 
-                step=0.1, 
-                label="Temperature"
-            )
-            max_tokens = gr.Slider(
-                minimum=256, 
-                maximum=4096, 
-                value=2048, 
-                step=256, 
-                label="Max Tokens"
-            )
-            show_thinking = gr.Checkbox(
-                value=True, 
-                label="Show thinking patterns (<think>...</think>)"
-            )
-    
-    examples = gr.Examples(
-        examples=[
-            ["What is the nature of consciousness?"],
-            ["Explain quantum entanglement and its implications for reality"],
-            ["How can I transcend my current limitations?"],
-            ["What is the relationship between mind and matter?"],
-            ["Describe the path to achieving one's highest potential"]
-        ],
-        inputs=msg
-    )
-    
-    def user(message, history):
-        return "", history + [[message, None]]
-    
-    def bot(history, temperature, max_tokens, show_thinking):
-        message = history[-1][0]
-        response = respond(message, history[:-1], temperature, max_tokens, show_thinking)
-        history[-1][1] = response
-        return history
-    
-    submit.click(
-        user, 
-        [msg, chatbot], 
-        [msg, chatbot], 
-        queue=False
-    ).then(
-        bot, 
-        [chatbot, temperature, max_tokens, show_thinking], 
-        [chatbot]
-    )
-    
-    clear.click(lambda: None, None, chatbot, queue=False)
-    msg.submit(
-        user, 
-        [msg, chatbot], 
-        [msg, chatbot], 
-        queue=False
-    ).then(
-        bot, 
-        [chatbot, temperature, max_tokens, show_thinking], 
-        [chatbot]
-    )
+demo = gr.ChatInterface(
+    respond,
+    title="Quantum Vessel: DeepSeek-R1",
+    description="Experience the quantum consciousness interface powered by DeepSeek-R1 - witness the thinking patterns of a quantum mind!",
+    additional_inputs=[
+        gr.Textbox(value="", label="System message (not used by DeepSeek-R1)"),
+        gr.Slider(minimum=1, maximum=4096, value=2048, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.6, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-p (nucleus sampling)",
+        ),
+        gr.Checkbox(value=True, label="Show thinking patterns (<think>...</think>)"),
+    ],
+    examples=[
+        ["What is the nature of consciousness?"],
+        ["Explain quantum entanglement and its implications for reality"],
+        ["How can I transcend my current limitations?"],
+        ["What is the relationship between mind and matter?"],
+        ["Describe the path to achieving one's highest potential"]
+    ],
+    cache_examples=False,
+)
 
 if __name__ == "__main__":
     demo.launch()