Update app.py

app.py

@@ -1,189 +1,146 @@
 import gradio as gr
-import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 import torch
 from threading import Thread
-import bitsandbytes as bnb
 
+# Model and device configuration
 phi4_model_path = "Compumacy/OpenBioLLm-70B"
-
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-# Load model with 4-bit quantization
-from transformers import BitsAndBytesConfig
+# === INITIALIZE EMPTY WEIGHTS ===
+init_empty_weights()
 
-# Configure 4-bit quantization
-quantization_config = BitsAndBytesConfig(
+# === CONFIGURE 4-BIT QUANTIZATION ===
+bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16,
     bnb_4bit_use_double_quant=True,
     bnb_4bit_quant_type="nf4"
 )
 
-phi4_model = AutoModelForCausalLM.from_pretrained(
+# === LOAD MODEL WITH QUANTIZATION ===
+model = AutoModelForCausalLM.from_pretrained(
+    phi4_model_path,
+    quantization_config=bnb_config,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+
+tokenizer = AutoTokenizer.from_pretrained(phi4_model_path)
+
+# === OFFLOAD TO CPU/DISK ===
+model = load_checkpoint_and_dispatch(
+    model,
     phi4_model_path,
     device_map="auto",
-    quantization_config=quantization_config
+    offload_folder="offload",
+    offload_state_dict=True,
+    max_memory={**{i: "12GB" for i in range(torch.cuda.device_count())}, "cpu": "30GB"}
 )
-phi4_tokenizer = AutoTokenizer.from_pretrained(phi4_model_path)
 
-@spaces.GPU(duration=120)
+# Enable gradient checkpointing if ever fine-tuning
+model.gradient_checkpointing_enable()
+
+# Optionally compile for PyTorch >= 2.0
+try:
+    model = torch.compile(model)
+except Exception:
+    pass
+
+# === RESPONSE GENERATOR ===
 def generate_response(user_message, max_tokens, temperature, top_k, top_p, repetition_penalty, history_state):
     if not user_message.strip():
         return history_state, history_state
-        
-    # Phi-4 model settings
-    model = phi4_model
-    tokenizer = phi4_tokenizer
-    start_tag = "<|im_start|>"
-    sep_tag = "<|im_sep|>"
-    end_tag = "<|im_end|>"
-
-    # Recommended prompt settings by Microsoft
-    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 setup
+    system_message = (
+        "Your role as an assistant involves thoroughly exploring questions through a systematic thinking process..."
+    )
+    start_tag, sep_tag, end_tag = "<|im_start|>", "<|im_sep|>", "<|im_end|>"
     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}"
+    for msg in history_state:
+        tag = msg["role"]
+        content = msg["content"]
+        prompt += f"{start_tag}{tag}{sep_tag}{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)
-
+    # Streaming setup
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-
-    # sampling techniques
     generation_kwargs = {
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
+        "input_ids": inputs.input_ids,
+        "attention_mask": inputs.attention_mask,
         "max_new_tokens": int(max_tokens),
         "do_sample": True,
-        "temperature": temperature,  # Use the slider value
+        "temperature": temperature,
         "top_k": int(top_k),
-        "top_p": top_p,              # Use the slider value
+        "top_p": top_p,
         "repetition_penalty": repetition_penalty,
-        "streamer": streamer,
+        "streamer": streamer
     }
 
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
+    # Run generation in thread
+    Thread(target=model.generate, kwargs=generation_kwargs).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()
+
+    # Stream tokens
+    for token in streamer:
+        clean = token.replace(start_tag, "").replace(sep_tag, "").replace(end_tag, "")
+        assistant_response += clean
+        new_history[-1]["content"] = assistant_response
         yield new_history, new_history
 
     yield new_history, new_history
 
+# === EXAMPLE MESSAGES ===
 example_messages = {
-    "Math reasoning": "If a rectangular prism has a length of 6 cm, a width of 4 cm, and a height of 5 cm, what is the length of the longest line segment that can be drawn from one vertex to another?",
-    "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?",
-    "Physics problem": "A ball is thrown upward with an initial velocity of 15 m/s from a height of 2 meters above the ground. Assuming the acceleration due to gravity is 9.8 m/s², determine: 1) The maximum height the ball reaches. 2) The total time the ball is in the air before hitting the ground. 3) The velocity with which the ball hits the ground."
+    "Math reasoning": "If a rectangular prism has a length of 6 cm...",
+    "Logic puzzle": "Four people (Alex, Blake, Casey, ...)",
+    "Physics problem": "A ball is thrown upward with an initial velocity..."
 }
 
+# === GRADIO APP ===
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # try the example problems below to see how the model breaks down complex reasoning problems.
-        ## *Running with 4-bit quantization*
-        """
-    )
-    
-    history_state = gr.State([])
+    gr.Markdown("""
+    # Phi-4 Chat
+    Try the example problems below to see how the model breaks down complex reasoning.
+    """ )
 
+    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"
-            )
+            max_tokens_slider = gr.Slider(64, 32768, step=1024, value=2048, 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"
-                )
-        
+                temperature_slider = gr.Slider(0.1, 2.0, value=0.8, label="Temperature")
+                top_k_slider = gr.Slider(1, 100, step=1, value=50, label="Top-k")
+                top_p_slider = gr.Slider(0.1, 1.0, value=0.95, label="Top-p")
+                repetition_penalty_slider = gr.Slider(1.0, 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
-                )
+                user_input = gr.Textbox(placeholder="Type your message...", 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("Physics problem")
+                for name in example_messages:
+                    btn = gr.Button(name)
+                    btn.click(fn=lambda n=name: gr.update(value=example_messages[n]), inputs=None, outputs=user_input)
 
     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]
-    ).then(
-        fn=lambda: gr.update(value=""),
-        inputs=None,
-        outputs=user_input
-    )
+    ).then(lambda: gr.update(value=""), None, user_input)
 
-    clear_button.click(
-        fn=lambda: ([], []),
-        inputs=None,
-        outputs=[chatbot, history_state]
-    )
-
-    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["Physics problem"]),
-        inputs=None,
-        outputs=user_input
-    )
+    clear_button.click(lambda: ([], []), None, [chatbot, history_state])
 
-demo.launch(ssr_mode=False)
+demo.launch(ssr_mode=False)