Update train_lora.py

train_lora.py

@@ -1,50 +1,180 @@
-import os, json, torch, wandb
-from transformers import (AutoModelForCausalLM, AutoTokenizer, Trainer,
-                          TrainingArguments, DataCollatorForLanguageModeling)
-from datasets import load_dataset
-from peft import LoraConfig, get_peft_model
-from huggingface_hub import HfApi
-
-HF_TOKEN = os.getenv("HF_TOKEN")
-WANDB_API_KEY = os.getenv("WANDB_API_KEY")
-wandb.login(key=WANDB_API_KEY)
-
-model_name = "deepseek-ai/deepseek-coder-1.3b-base"
-dataset = load_dataset("westenfelder/NL2SH-ALFA")
-
-tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-def tok_fn(b): return tok([f"{n} => {bsh}" for n,bsh in zip(b['nl'],b['bash'])],
-                          truncation=True,padding="max_length",max_length=512)
-train, test = dataset["train"].map(tok_fn,batched=True), dataset["test"].map(tok_fn,batched=True)
-
-m = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16,
-                                         low_cpu_mem_usage=True, device_map="auto",
-                                         trust_remote_code=True)
-m.config.use_cache=False
-for p in m.parameters(): p.requires_grad=False
-
-cfg=LoraConfig(r=8,lora_alpha=16,target_modules=["q_proj","v_proj","k_proj","o_proj",
-                                                 "gate_proj","down_proj","up_proj"],
-               lora_dropout=0.05,bias="none",task_type="CAUSAL_LM")
-m=get_peft_model(m,cfg)
-coll=DataCollatorForLanguageModeling(tokenizer=tok,mlm=False)
-
-args=TrainingArguments(output_dir="./out",num_train_epochs=1,per_device_train_batch_size=1,
-                       gradient_accumulation_steps=8,learning_rate=2e-4,fp16=True,
-                       save_strategy="epoch",logging_steps=25,report_to=["wandb"])
-t=Trainer(model=m,args=args,train_dataset=train,eval_dataset=test,data_collator=coll)
-wandb.init(project="deepseek-qlora-monthly",name="deepseek-lite-run")
-t.train()
-
-metrics=t.evaluate(); acc=1-metrics.get("eval_loss",1)
-with open("out/metrics.json","w") as f: json.dump(metrics,f)
-wandb.log({"accuracy":acc})
-print(f"✅ Eval accuracy {acc:.4f}")
-
-ad="out/lora_adapters"; os.makedirs(ad,exist_ok=True)
-m.save_pretrained(ad); tok.save_pretrained(ad)
-artifact=wandb.Artifact("deepseek-lora-adapters","model"); artifact.add_dir(ad); wandb.log_artifact(artifact)
-
-api=HfApi(token=HF_TOKEN)
-api.upload_folder(folder_path=ad,repo_id="your-username/deepseek-lora-monthly",path_in_repo=".")
-print("✅ Uploaded to HF Hub")
+#!/usr/bin/env python3
+"""
+train_lora.py
+- Fine-tune DeepSeek 1.3B with LoRA (QLoRA-ish setup)
+- Save adapters using safe_serialization=True -> adapter_model.safetensors
+- Upload adapter folder to Hugging Face Hub (VaibhavHD/deepseek-lora-monthly)
+- Log metrics/artifact to Weights & Biases
+"""
+
+import os
+import json
+import wandb
+import torch
+from huggingface_hub import HfApi
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM,
+    TrainingArguments, Trainer, DataCollatorForLanguageModeling
+)
+from peft import LoraConfig, get_peft_model
+
+# -----------------------------
+# Config (edit if needed)
+# -----------------------------
+HF_REPO = "VaibhavHD/deepseek-lora-monthly"   # your HF model repo
+MODEL_NAME = "deepseek-ai/deepseek-coder-1.3b-base"
+OUT_DIR = "out"
+ADAPTER_DIR = os.path.join(OUT_DIR, "lora_adapters")
+
+# env secrets expected:
+HF_TOKEN = os.getenv("HF_TOKEN")
+WANDB_API_KEY = os.getenv("WANDB_API_KEY")
+
+if WANDB_API_KEY:
+    wandb.login(key=WANDB_API_KEY)
+else:
+    print("⚠️ WANDB_API_KEY not found in env; continuing without W&B logging.")
+
+# -----------------------------
+# Load dataset
+# -----------------------------
+print("Loading dataset...")
+dataset = load_dataset("westenfelder/NL2SH-ALFA")
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+
+def tokenize_fn(batch):
+    texts = [f"{nl} => {bash}" for nl, bash in zip(batch["nl"], batch["bash"])]
+    return tokenizer(texts, truncation=True, padding="max_length", max_length=512)
+
+train = dataset["train"].map(tokenize_fn, batched=True)
+test  = dataset["test"].map(tokenize_fn, batched=True)
+
+# Optional small-subset for fast runs (uncomment to use)
+# train = train.shuffle(seed=42).select(range(200))
+# test  = test.shuffle(seed=42).select(range(20))
+
+# -----------------------------
+# Load base model (half precision)
+# -----------------------------
+print("Loading base model (may take a moment)...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+# avoid caching issues
+model.config.use_cache = False
+for p in model.parameters():
+    p.requires_grad = False
+
+# -----------------------------
+# Attach LoRA
+# -----------------------------
+print("Attaching LoRA adapters...")
+lora_config = LoraConfig(
+    r=8,
+    lora_alpha=16,
+    target_modules=[
+        "q_proj", "v_proj", "k_proj", "o_proj",
+        "gate_proj", "down_proj", "up_proj"
+    ],
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM",
+)
+model = get_peft_model(model, lora_config)
+
+# -----------------------------
+# Data collator + training args
+# -----------------------------
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+training_args = TrainingArguments(
+    output_dir=OUT_DIR,
+    num_train_epochs=1,
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=8,
+    learning_rate=2e-4,
+    fp16=True,
+    save_strategy="epoch",
+    logging_steps=25,
+    report_to=["wandb"] if WANDB_API_KEY else [],
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train,
+    eval_dataset=test,
+    data_collator=data_collator,
+)
+
+# -----------------------------
+# Run training
+# -----------------------------
+print("Starting training...")
+if WANDB_API_KEY:
+    wandb.init(project="deepseek-qlora-monthly", name="deepseek-lite-run")
+
+trainer.train()
+
+# -----------------------------
+# Evaluate and save metrics
+# -----------------------------
+print("Evaluating...")
+metrics = trainer.evaluate()
+# compute simple "accuracy-like" metric from loss (replace with real metric if you have one)
+new_acc = 1.0 - metrics.get("eval_loss", 1.0)
+print(f"Eval metrics: {metrics}")
+print(f"Pseudo-accuracy (1 - eval_loss): {new_acc:.6f}")
+
+os.makedirs(ADAPTER_DIR, exist_ok=True)
+metrics_path = os.path.join(OUT_DIR, "metrics.json")
+with open(metrics_path, "w") as f:
+    json.dump(metrics, f)
+
+if WANDB_API_KEY:
+    wandb.log({"accuracy": new_acc})
+    # log artifact
+    artifact = wandb.Artifact(
+        name="deepseek-lora-adapters",
+        type="model",
+        description="LoRA adapters saved with safe_serialization"
+    )
+
+# -----------------------------
+# Save adapters using safe_serialization
+# -----------------------------
+print("Saving adapters with safe_serialization=True (produces .safetensors)...")
+model.save_pretrained(ADAPTER_DIR, safe_serialization=True)
+tokenizer.save_pretrained(ADAPTER_DIR)
+
+# add to wandb artifact directory
+if WANDB_API_KEY:
+    artifact.add_dir(ADAPTER_DIR)
+    wandb.log_artifact(artifact, aliases=["latest"])
+
+print(f"Adapters saved to: {ADAPTER_DIR}")
+print("Files in adapter dir:", os.listdir(ADAPTER_DIR))
+
+# -----------------------------
+# Upload to Hugging Face model repo
+# -----------------------------
+if HF_TOKEN:
+    print(f"Uploading adapter folder to Hugging Face repo: {HF_REPO}")
+    api = HfApi()
+    # upload_folder will overwrite same filenames in the repo
+    api.upload_folder(
+        folder_path=ADAPTER_DIR,
+        path_in_repo=".",
+        repo_id=HF_REPO,
+        token=HF_TOKEN
+    )
+    print("✅ Upload complete.")
+else:
+    print("⚠️ HF_TOKEN not set. Skipping upload to Hugging Face Hub.")