Update train.py

train.py

@@ -1,8 +1,8 @@
 """
 DocuMint Smart Training Pipeline
-- Core Adapter (Frozen)
-- Skill-wise Adapters
-- Safe Continual Learning
+- Core adapter (one-time training)
+- Skill-wise adapters (additive learning)
+- Safe continual learning (no destruction)
 """
 
 import os
@@ -29,8 +29,9 @@ from huggingface_hub import login
 
 BASE_MODEL = "Qwen/Qwen2-0.5B-Instruct"
 
-CORE_ADAPTER_REPO = "himu1780/DocuMint-Core"
-SKILL_ADAPTER_REPO = "himu1780/DocuMint-Skill"   # will be suffixed
+CORE_REPO = "himu1780/DocuMint-Core"
+SKILL_REPO_PREFIX = "himu1780/DocuMint-Skill"
+
 OUTPUT_DIR = "./lora_output"
 
 MAX_LENGTH = 512
@@ -77,28 +78,29 @@ def format_example(ex):
         )
     else:
         text = ex.get("text", str(ex))
+
     return {"text": text}
 
 
 def prepare_dataset(tokenizer, dataset_name):
-    ds = load_dataset(dataset_name, split="train")
-    ds = ds.map(format_example, remove_columns=ds.column_names)
+    dataset = load_dataset(dataset_name, split="train")
+    dataset = dataset.map(format_example, remove_columns=dataset.column_names)
 
     def tokenize(ex):
-        out = tokenizer(
+        tokens = tokenizer(
             ex["text"],
             truncation=True,
             padding="max_length",
             max_length=MAX_LENGTH,
         )
-        out["labels"] = out["input_ids"].copy()
-        return out
+        tokens["labels"] = tokens["input_ids"].copy()
+        return tokens
 
-    ds = ds.map(tokenize, remove_columns=["text"])
-    return ds
+    dataset = dataset.map(tokenize, remove_columns=["text"])
+    return dataset
 
 
-# ================== MODEL LOADING ==================
+# ================== MODEL ==================
 
 def load_base():
     tokenizer = AutoTokenizer.from_pretrained(
@@ -109,7 +111,7 @@ def load_base():
 
     model = AutoModelForCausalLM.from_pretrained(
         BASE_MODEL,
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float32,   # CPU safe
         device_map="cpu",
         trust_remote_code=True,
         low_cpu_mem_usage=True,
@@ -117,7 +119,7 @@ def load_base():
     return model, tokenizer
 
 
-def build_lora():
+def lora_config():
     return LoraConfig(
         r=8,
         lora_alpha=16,
@@ -128,37 +130,35 @@ def build_lora():
     )
 
 
-# ================== CORE + SKILL LOGIC ==================
+# ================== ADAPTER LOGIC ==================
 
 def load_core_adapter(model):
-    """
-    Core adapter is ALWAYS frozen.
-    """
-    if not os.path.exists(OUTPUT_DIR + "/core"):
+    core_path = os.path.join(OUTPUT_DIR, "core")
+
+    if not os.path.exists(core_path):
         raise RuntimeError("Core adapter not found. Train core first.")
 
-    model = PeftModel.from_pretrained(model, OUTPUT_DIR + "/core")
+    model = PeftModel.from_pretrained(model, core_path)
+
+    # Freeze everything
     for p in model.parameters():
         p.requires_grad = False
 
-    print("🧠 Core adapter loaded & frozen")
+    print("🧠 Core adapter loaded and frozen")
     return model
 
 
-def load_or_create_skill_adapter(model, skill_name):
-    """
-    Skill adapter is trainable.
-    """
-    skill_path = f"{OUTPUT_DIR}/{skill_name}"
+def load_or_create_adapter(model, skill_name):
+    adapter_path = os.path.join(OUTPUT_DIR, skill_name)
 
-    if os.path.exists(skill_path):
-        print(f"🔁 Loading existing skill adapter: {skill_name}")
+    if os.path.exists(adapter_path):
+        print(f"🔁 Loading existing adapter: {skill_name}")
         model = PeftModel.from_pretrained(
-            model, skill_path, is_trainable=True
+            model, adapter_path, is_trainable=True
         )
     else:
-        print(f"🆕 Creating new skill adapter: {skill_name}")
-        model = get_peft_model(model, build_lora())
+        print(f"🆕 Creating new adapter: {skill_name}")
+        model = get_peft_model(model, lora_config())
 
     model.print_trainable_parameters()
     return model
@@ -174,18 +174,22 @@ def train_skill(
     batch_size: int,
 ):
     """
-    Train a skill adapter WITHOUT destroying core knowledge.
+    skill_name:
+    - "core"  -> core training (one time)
+    - others  -> skill training (requires core)
     """
 
     hf_auth()
 
     model, tokenizer = load_base()
 
-    # Load & freeze core
-    model = load_core_adapter(model)
+    # IMPORTANT FIX:
+    # Load core ONLY if training a skill
+    if skill_name != "core":
+        model = load_core_adapter(model)
 
-    # Load / create skill adapter
-    model = load_or_create_skill_adapter(model, skill_name)
+    # Load or create adapter
+    model = load_or_create_adapter(model, skill_name)
 
     dataset = prepare_dataset(tokenizer, dataset_name)
 
@@ -206,7 +210,8 @@ def train_skill(
     )
 
     collator = DataCollatorForLanguageModeling(
-        tokenizer=tokenizer, mlm=False
+        tokenizer=tokenizer,
+        mlm=False,
     )
 
     trainer = Trainer(
@@ -218,31 +223,32 @@ def train_skill(
 
     trainer.train()
 
-    # Save skill adapter
-    save_path = f"{OUTPUT_DIR}/{skill_name}"
+    # Save locally
+    save_path = os.path.join(OUTPUT_DIR, skill_name)
     model.save_pretrained(save_path)
     tokenizer.save_pretrained(save_path)
 
-    # Push to HF
-    model.push_to_hub(f"{SKILL_ADAPTER_REPO}-{skill_name}")
-    tokenizer.push_to_hub(f"{SKILL_ADAPTER_REPO}-{skill_name}")
+    # Push to Hub
+    if skill_name == "core":
+        repo = CORE_REPO
+    else:
+        repo = f"{SKILL_REPO_PREFIX}-{skill_name}"
+
+    model.push_to_hub(repo)
+    tokenizer.push_to_hub(repo)
 
     cleanup()
-    print(f"✅ Skill '{skill_name}' trained and saved")
+    print(f"✅ Training finished for adapter: {skill_name}")
 
 
-# ================== ROUTING (INFERENCE READY) ==================
+# ================== ROUTING (INFERENCE) ==================
 
 def load_for_inference(skill_name: str):
-    """
-    Adapter routing:
-    Base → Core → Skill
-    """
     model, tokenizer = load_base()
 
-    model = PeftModel.from_pretrained(model, CORE_ADAPTER_REPO)
+    model = PeftModel.from_pretrained(model, CORE_REPO)
     model = PeftModel.from_pretrained(
-        model, f"{SKILL_ADAPTER_REPO}-{skill_name}"
+        model, f"{SKILL_REPO_PREFIX}-{skill_name}"
     )
 
     model.eval()
@@ -254,4 +260,4 @@ def load_for_inference(skill_name: str):
 
 if __name__ == "__main__":
     print("🏆 DocuMint Smart Training System Ready")
-    print("Use train_skill() to add new intelligence safely")
+    print("Use train_skill() to train core or add skills safely")