Update app.py

app.py

@@ -8,173 +8,68 @@ veri_model_path = "nyu-dice-lab/VeriThoughts-Reasoning-7B"
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-veri_model = AutoModelForCausalLM.from_pretrained(veri_model_path, device_map="auto", torch_dtype="auto")
-veri_tokenizer = AutoTokenizer.from_pretrained(veri_model_path)
+# Try loading the model with explicit error handling
+try:
+    veri_model = AutoModelForCausalLM.from_pretrained(veri_model_path, device_map="auto", torch_dtype="auto")
+    veri_tokenizer = AutoTokenizer.from_pretrained(veri_model_path)
+except Exception as e:
+    print(f"Model loading error: {e}")
+    veri_model = None
+    veri_tokenizer = None
 
 @spaces.GPU(duration=60)
-def generate_response(user_message, max_tokens, temperature, top_k, top_p, repetition_penalty, history_state):
+def generate_response(user_message, history):
+    if not veri_model or not veri_tokenizer:
+        return history + [["Error", "Model not loaded properly"]]
+    
     if not user_message.strip():
-        return history_state, history_state
+        return history
         
-    # model settings
-    model = veri_model
-    tokenizer = veri_tokenizer
-    start_tag = "<|im_start|>"
-    sep_tag = "<|im_sep|>"
-    end_tag = "<|im_end|>"
-
-    # Recommended prompt settings by Qwen
-    system_message = "Your role as an assistant involves thoroughly exploring questions through a systematic thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracing, and iteration to develop well-considered thinking process. Please structure your response into two main sections: Thought and Solution using the specified format: <think> {Thought section} </think> {Solution section}. In the Thought section, detail your reasoning process in steps. Each step should include detailed considerations such as analysing questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The Solution section should be logical, accurate, and concise and detail necessary steps needed to reach the conclusion. Now, try to solve the following question through the above guidelines:"
-    prompt = f"{start_tag}system{sep_tag}{system_message}{end_tag}"
-    for message in history_state:
-        if message["role"] == "user":
-            prompt += f"{start_tag}user{sep_tag}{message['content']}{end_tag}"
-        elif message["role"] == "assistant" and message["content"]:
-            prompt += f"{start_tag}assistant{sep_tag}{message['content']}{end_tag}"
-    prompt += f"{start_tag}user{sep_tag}{user_message}{end_tag}{start_tag}assistant{sep_tag}"
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-
-    do_sample = not (temperature == 1.0 and top_k >= 100 and top_p == 1.0)
-
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-
-    # sampling techniques
-    generation_kwargs = {
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
-        "max_new_tokens": int(max_tokens),
-        "do_sample": True,
-        "temperature": 0.8,
-        "top_k": int(top_k),
-        "top_p": 0.95,
-        "repetition_penalty": repetition_penalty,
-        "streamer": streamer,
-    }
-
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    # Stream the response
-    assistant_response = ""
-    new_history = history_state + [
-        {"role": "user", "content": user_message},
-        {"role": "assistant", "content": ""}
-    ]
-    for new_token in streamer:
-        cleaned_token = new_token.replace("<|im_start|>", "").replace("<|im_sep|>", "").replace("<|im_end|>", "")
-        assistant_response += cleaned_token
-        new_history[-1]["content"] = assistant_response.strip()
-        yield new_history, new_history
-
-    yield new_history, new_history
-
-# Fixed: Match the keys with your button labels
-example_messages = {
-    "Math reasoning": "Calculate the sum of the first 10 prime numbers and explain your reasoning step by step.",
-    "Logic puzzle": "Four people (Alex, Blake, Casey, and Dana) each have a different favorite color (red, blue, green, yellow) and a different favorite fruit (apple, banana, cherry, date). Given the following clues: 1) The person who likes red doesn't like dates. 2) Alex likes yellow. 3) The person who likes blue likes cherries. 4) Blake doesn't like apples or bananas. 5) Casey doesn't like yellow or green. Who likes what color and what fruit?",
-    "Verilog example": "Design a 4-bit adder circuit in Verilog with proper test benches."
-}
-
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # VeriThoughts-7B Chatbot 
-        Welcome to VeriThoughts-7B! This is a reasoning model for Verilog code generation. 
-        
-        The model will provide responses with two sections:
-        1. **<think>**: A detailed step-by-step reasoning process showing its work
-        2. **Solution**: A concise, accurate final answer based on the reasoning
-        
-        """
-    )
+    # Simple generation without streaming first
+    system_message = "You are a helpful assistant that thinks step by step."
     
-    history_state = gr.State([])
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### Settings")
-            max_tokens_slider = gr.Slider(
-                minimum=64,
-                maximum=32768,
-                step=1024,
-                value=4096,
-                label="Max Tokens"
-            )
-            with gr.Accordion("Advanced Settings", open=False):
-                temperature_slider = gr.Slider(
-                    minimum=0.1,
-                    maximum=2.0,
-                    value=0.8,
-                    label="Temperature"
-                )
-                top_k_slider = gr.Slider(
-                    minimum=1,
-                    maximum=100,
-                    step=1,
-                    value=50,
-                    label="Top-k"
-                )
-                top_p_slider = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=0.95,
-                    label="Top-p"
-                )
-                repetition_penalty_slider = gr.Slider(
-                    minimum=1.0,
-                    maximum=2.0,
-                    value=1.0,
-                    label="Repetition Penalty"
-                )
-        
-        with gr.Column(scale=4):
-            chatbot = gr.Chatbot(label="Chat", type="messages")
-            with gr.Row():
-                user_input = gr.Textbox(
-                    label="Your message",
-                    placeholder="Type your message here...",
-                    scale=3
-                )
-                submit_button = gr.Button("Send", variant="primary", scale=1)
-                clear_button = gr.Button("Clear", scale=1)
-            gr.Markdown("**Try these examples:**")
-            with gr.Row():
-                example1_button = gr.Button("Math reasoning")
-                example2_button = gr.Button("Logic puzzle")
-                example3_button = gr.Button("Verilog example")
+    # Create conversation history
+    conversation = f"System: {system_message}\n"
+    for h in history:
+        conversation += f"User: {h[0]}\nAssistant: {h[1]}\n"
+    conversation += f"User: {user_message}\nAssistant:"
+    
+    inputs = veri_tokenizer(conversation, return_tensors="pt", truncation=True, max_length=2048).to(device)
+    
+    with torch.no_grad():
+        outputs = veri_model.generate(
+            **inputs,
+            max_new_tokens=512,
+            temperature=0.8,
+            do_sample=True,
+            pad_token_id=veri_tokenizer.eos_token_id
+        )
+    
+    response = veri_tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    
+    # Return updated history
+    return history + [[user_message, response.strip()]]
 
-    submit_button.click(
-        fn=generate_response,
-        inputs=[user_input, max_tokens_slider, temperature_slider, top_k_slider, top_p_slider, repetition_penalty_slider, history_state],
-        outputs=[chatbot, history_state]
+# Create minimal interface
+with gr.Blocks(title="VeriThoughts-7B Chatbot") as demo:
+    gr.Markdown("# VeriThoughts-7B Chatbot")
+    
+    chatbot = gr.Chatbot(value=[], label="Chat")
+    msg = gr.Textbox(label="Your message", placeholder="Type here...")
+    clear = gr.Button("Clear")
+    
+    # Simple event handling
+    msg.submit(
+        fn=generate_response, 
+        inputs=[msg, chatbot], 
+        outputs=chatbot
     ).then(
-        fn=lambda: gr.update(value=""),
-        inputs=None,
-        outputs=user_input
-    )
-
-    clear_button.click(
-        fn=lambda: ([], []),
-        inputs=None,
-        outputs=[chatbot, history_state]
-    )
-
-    # Fixed: Now these will work without KeyError
-    example1_button.click(
-        fn=lambda: gr.update(value=example_messages["Math reasoning"]),
-        inputs=None,
-        outputs=user_input
-    )
-    example2_button.click(
-        fn=lambda: gr.update(value=example_messages["Logic puzzle"]),
-        inputs=None,
-        outputs=user_input
-    )
-    example3_button.click(
-        fn=lambda: gr.update(value=example_messages["Verilog example"]),
-        inputs=None,
-        outputs=user_input
+        lambda: "", 
+        inputs=None, 
+        outputs=msg
     )
+    
+    clear.click(lambda: [], outputs=chatbot)
 
-demo.launch(ssr_mode=False, share=True)
+# Launch without ssr_mode parameter which might cause issues
+demo.launch()