Update app.py

app.py

@@ -18,10 +18,42 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def check_gpu_status():
     return f"✅ GPU: {torch.cuda.get_device_name(0)}" if device == "cuda" else "⚠️ Using CPU only"
+# ------------------------------------------------------
+# 🧩 Download Dataset to /tmp/
+# ------------------------------------------------------
+def download_gita_dataset():
+    repo_id = "rahul7star/Gita"
+    local_dir = "/tmp/gita_data"
+
+    if os.path.exists(local_dir):
+        shutil.rmtree(local_dir)
+    os.makedirs(local_dir, exist_ok=True)
+
+    print(f"📥 Downloading dataset from {repo_id} ...")
+    snapshot_download(repo_id=repo_id, local_dir=local_dir, repo_type="dataset")
+
+    # Try to locate the CSV file
+    csv_path = None
+    for root, _, files in os.walk(local_dir):
+        for f in files:
+            if f.lower().endswith(".csv"):
+                csv_path = os.path.join(root, f)
+                break
+    if not csv_path:
+        raise FileNotFoundError("No CSV file found in the Gita dataset repository.")
+
+    print(f"✅ Found CSV: {csv_path}")
+    return csv_path
+
+
+# ------------------------------------------------------
+# 🚀 Training function
+# ------------------------------------------------------
 
 # ---------------------------------------------------------------------
 # Training Logic
 # ---------------------------------------------------------------------
+@spaces.GPU(duration=300)
 def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress(track_tqdm=True)):
     output_log = []