Update app_low.py

app_low.py

@@ -1,18 +1,109 @@
-# Load model directly
+import os
+import torch
+import gradio as gr
+from huggingface_hub import snapshot_download
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-tokenizer = AutoTokenizer.from_pretrained("rubricreward/mR3-Qwen3-14B-en-prompt-en-thinking")
-model = AutoModelForCausalLM.from_pretrained("rubricreward/mR3-Qwen3-14B-en-prompt-en-thinking")
-messages = [
-    {"role": "user", "content": "Who are you?"},
-]
-inputs = tokenizer.apply_chat_template(
-	messages,
-	add_generation_prompt=True,
-	tokenize=True,
-	return_dict=True,
-	return_tensors="pt",
-).to(model.device)
-
-outputs = model.generate(**inputs, max_new_tokens=40)
-print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))
+# =========================================================
+# 1️⃣ Download model from HF (once)
+# =========================================================
+MODEL_ID = "rubricreward/mR3-Qwen3-14B-en-prompt-en-thinking"
+print(f"📦 Downloading or loading cached model: {MODEL_ID} ...")
+model_path = snapshot_download(repo_id=MODEL_ID)
+print(f"✅ Model path: {model_path}")
+
+# =========================================================
+# 2️⃣ Smart device setup (auto CPU/GPU/offload)
+# =========================================================
+if torch.cuda.is_available():
+    device = "cuda"
+    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    print("⚙️ Using CUDA for inference.")
+else:
+    device = "cpu"
+    dtype = torch.float32
+    print("⚙️ Using CPU — applying efficient offloading settings.")
+
+# =========================================================
+# 3️⃣ Load model and tokenizer with optimized config
+# =========================================================
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    torch_dtype=dtype,
+    device_map="auto" if device == "cuda" else {"": "cpu"},
+    low_cpu_mem_usage=True,
+    offload_folder="./offload" if device == "cpu" else None,
+)
+
+# =========================================================
+# 4️⃣ Chat inference function
+# =========================================================
+def chat_with_model(user_input, chat_history):
+    """Run chat-style inference."""
+    if not user_input.strip():
+        return chat_history + [["", "⚠️ Please enter a message."]]
+
+    messages = [{"role": "user", "content": user_input}]
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_tensors="pt",
+    ).to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            repetition_penalty=1.1,
+        )
+
+    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    chat_history = chat_history + [[user_input, response.strip()]]
+    return chat_history
+
+# =========================================================
+# 5️⃣ Gradio Chat UI
+# =========================================================
+with gr.Blocks(theme=gr.themes.Soft(), title="💬 Qwen3-14B Thinking Chat") as demo:
+    gr.Markdown(
+        """
+        # 🧠 mR3-Qwen3-14B Prompt-Enhanced Chat  
+        A reasoning-ready English chat model from **RubricReward**, optimized for CPU/GPU with offloading.  
+        ---
+        """
+    )
+
+    chatbot = gr.Chatbot(height=400, label="Chat with mR3-Qwen3-14B")
+    user_input = gr.Textbox(placeholder="Type your message here...", label="Your Message", lines=2)
+    send_button = gr.Button("🚀 Send", variant="primary")
+
+    def clear_history():
+        return []
+
+    clear_btn = gr.Button("🧹 Clear Chat")
+
+    send_button.click(chat_with_model, [user_input, chatbot], chatbot)
+    user_input.submit(chat_with_model, [user_input, chatbot], chatbot)
+    clear_btn.click(fn=clear_history, outputs=chatbot)
+
+    gr.Markdown(
+        """
+        ---
+        💡 **Tips:**  
+        - Ask it to explain, reason, or summarize complex ideas.  
+        - Try: *“Explain quantum computing in simple terms.”*  
+        - Try: *“Give me 3 creative ways to promote a science exhibition.”*
+        """
+    )
+
+# =========================================================
+# 6️⃣ Launch App
+# =========================================================
+if __name__ == "__main__":
+    demo.launch()