Update app.py

app.py

@@ -6,6 +6,7 @@ import uuid
 import time
 import sys
 import gc
+import multiprocessing
 from datetime import datetime
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from itertools import chain
@@ -36,6 +37,7 @@ logging.basicConfig(level=logging.CRITICAL, stream=sys.stderr)
 if torch.cuda.is_available():
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
 
 JOBS = {}
 
@@ -107,7 +109,7 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
     
     job = JOBS[job_id]
     job.status = "RUNNING"
-    job.add_log("System: Starting Neural Forge Engine...")
+    job.add_log("System: Starting High-Performance Neural Engine...")
 
     try:
         if not hf_token.startswith("hf_"):
@@ -116,6 +118,7 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
         os.environ["WANDB_DISABLED"] = "true"
         os.environ["HF_TOKEN"] = hf_token
         os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+        os.environ["TOKENIZERS_PARALLELISM"] = "true"
         
         login(token=hf_token)
         try:
@@ -151,9 +154,10 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
                 return None
 
         streams = []
-        job.set_progress(0.05, "Data: Connecting streams...")
+        job.set_progress(0.05, "Data: Connecting streams (Multi-threaded)...")
         
-        with ThreadPoolExecutor(max_workers=4) as executor:
+        cpu_count = multiprocessing.cpu_count()
+        with ThreadPoolExecutor(max_workers=cpu_count * 2) as executor:
             futures = []
             for ds_name in dataset_list:
                 futures.append(executor.submit(load_single, ds_name, None))
@@ -166,7 +170,7 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
         if not streams:
             raise Exception("No valid datasets found")
 
-        job.set_progress(0.1, f"Data: {len(streams)} sources active.")
+        job.set_progress(0.1, f"Data: {len(streams)} high-speed sources active.")
 
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left", add_eos_token=True, add_bos_token=True)
         if tokenizer.pad_token is None:
@@ -180,16 +184,15 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
                     text = str(item.get("text", item.get("content", str(item))))
                     if len(text) < 10: continue
                     batch_buffer.append(text)
-                    if len(batch_buffer) >= 20:
-                        for txt in batch_buffer:
-                            tokens = tokenizer(txt, truncation=True, max_length=1024)
-                            tokens["labels"] = tokens["input_ids"].copy()
-                            yield tokens
+                    if len(batch_buffer) >= 100:
+                        encoded_batch = tokenizer(batch_buffer, truncation=True, max_length=2048, padding=False)
+                        for input_ids in encoded_batch["input_ids"]:
+                            yield {"input_ids": input_ids, "labels": input_ids}
                         batch_buffer = []
                 except:
                     continue
 
-        job.set_progress(0.15, "Model: Loading weights...")
+        job.set_progress(0.15, "Model: Loading weights (Fast IO)...")
         
         torch.cuda.empty_cache()
         gc.collect()
@@ -198,7 +201,8 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
             model_name, 
             trust_remote_code=True, 
             device_map="auto",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            low_cpu_mem_usage=True
         )
         
         peft_config = LoraConfig(
@@ -221,14 +225,17 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
             gradient_accumulation_steps=4,
             max_steps=int(train_steps),
             learning_rate=learning_rate,
-            optim="adamw_torch",
+            optim="adamw_torch_fused" if torch.cuda.is_available() else "adamw_torch",
             logging_steps=1,
             save_strategy="steps",
             save_steps=max(10, int(int(train_steps)/5)),
             save_total_limit=2,
             report_to="none",
             fp16=True if torch.cuda.is_available() else False,
-            disable_tqdm=True
+            disable_tqdm=True,
+            dataloader_num_workers=4,
+            dataloader_pin_memory=True,
+            torch_compile=True if os.name == 'posix' else False
         )
 
         dataset_iterable = IterableDataset.from_generator(process_stream_generator)
@@ -240,7 +247,7 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
             callbacks=[CustomTrainerCallback(job_id, hf_token, full_repo_id)]
         )
 
-        job.set_progress(0.2, "Training: Phase initiated...")
+        job.set_progress(0.2, "Training: Accelerated Phase initiated...")
         trainer.train()
         
         job.set_progress(0.9, "Processing: Merging tensors...")
@@ -254,7 +261,8 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
             return_dict=True,
             torch_dtype=torch.float16,
             trust_remote_code=True,
-            device_map="auto"
+            device_map="auto",
+            low_cpu_mem_usage=True
         )
         
         model_to_merge = PeftModel.from_pretrained(base_reload, output_dir)
@@ -337,7 +345,7 @@ def load_from_url(request: gr.Request):
         pass
     return gr.update(selected="launch_tab"), ""
 
-with gr.Blocks(title="Nucleus Enterprise") as demo:
+with gr.Blocks(title="Nucleus Enterprise", theme=gr.themes.Base()) as demo:
     with gr.Column():
         gr.Markdown("# ⚛️ NUCLEUS ENTERPRISE")
         gr.Markdown("Autonomous LLM Foundry | V5.0 Stable")