Update app.py

app.py

@@ -5,7 +5,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 """
 🧮 Root_Math full model chat app
-Auto-detects GPU/CPU and loads appropriate base model.
+Supports private/public repo and GPU/CPU auto-detection.
 """
 
 # ✅ Load Hugging Face API token
@@ -13,29 +13,23 @@ api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 if not api_token:
     raise ValueError("❌ ERROR: Hugging Face API token is not set. Please set it as an environment variable.")
 
-# ✅ Detect environment
-use_cuda = torch.cuda.is_available()
+# ✅ Correct model repo name (case-sensitive)
+model_name = "Hrushi02/Root_Math"  # double-check on HF website
 
-if use_cuda:
-    print("🚀 GPU detected — using float16 model for efficiency.")
-    dtype = torch.float16
-else:
-    print("💻 CPU detected — using float32 model.")
-    dtype = torch.float32
+# ✅ Device and dtype
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+print(f"⚡ Loading model on {device.upper()} with dtype={dtype}")
 
-# ✅ Load your full fine-tuned model directly
-model_name = "Hrushi02/Root_Math"  # Your repo
-
-print(f"🔄 Loading full model: {model_name} ...")
+# ✅ Load the model (directly, full fine-tuned)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=dtype,
     device_map="auto",
-    token=api_token
+    token=api_token  # required for private repo
 )
 
 # ✅ Load tokenizer
-print("🔄 Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(model_name, token=api_token)
 
 
@@ -68,7 +62,7 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
     yield response
 
 
-# ✅ Gradio UI
+# ✅ Gradio ChatInterface
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[