hftoken

fine_tune_llama32_1b.py

@@ -1,13 +1,99 @@
 import torch
+import os
+import json
+import urllib.request
+import urllib.error
+import getpass
 from datasets import load_dataset
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
 )
-from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model
+from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model, AutoPeftModelForCausalLM
 from trl import SFTTrainer, SFTConfig
 
+
+def _fetch_secret_json(url: str, api_key: str, timeout_s: int = 30) -> str:
+    req = urllib.request.Request(
+        url,
+        headers={
+            "accept": "application/json",
+            "X-API-Key": api_key,
+        },
+        method="GET",
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
+            raw = resp.read().decode("utf-8")
+    except urllib.error.HTTPError as e:
+        body = e.read().decode("utf-8", errors="replace") if hasattr(e, "read") else ""
+        raise RuntimeError(f"KeyVault HTTP {e.code}: {body}") from e
+    except Exception as e:
+        raise RuntimeError(f"KeyVault request failed: {e}") from e
+
+    try:
+        payload = json.loads(raw)
+    except json.JSONDecodeError as e:
+        raise RuntimeError(f"KeyVault did not return JSON: {raw[:200]}") from e
+
+    if isinstance(payload, str) and payload.strip():
+        return payload.strip()
+
+    if isinstance(payload, dict):
+        for key in ("value", "secret", "token", "hftoken", "hf_token"):
+            value = payload.get(key)
+            if isinstance(value, str) and value.strip():
+                return value.strip()
+        if len(payload) == 1:
+            value = next(iter(payload.values()))
+            if isinstance(value, str) and value.strip():
+                return value.strip()
+
+    raise RuntimeError(f"Unexpected KeyVault JSON shape: {type(payload).__name__}")
+
+
+def get_hf_token() -> str | None:
+    """Get a Hugging Face token without persisting it.
+
+    Order:
+    1) Use `HF_TOKEN` if set (runtime-only, provided by caller)
+    2) Else fetch from KeyVault URL (default: /secrets/hftoken) using an API key
+       - `KEYVAULT_API_KEY` env var, or
+       - prompt at runtime (getpass)
+
+    This intentionally avoids calling `huggingface_hub.login()` to prevent writing
+    tokens to disk.
+    """
+
+    token = os.environ.get("HF_TOKEN")
+    if token and token.strip():
+        return token.strip()
+
+    keyvault_url = os.environ.get(
+        "KEYVAULT_HF_TOKEN_URL",
+        "https://maxxcarl-keyvault.hf.space/secrets/hftoken",
+    )
+
+    api_key = os.environ.get("KEYVAULT_API_KEY")
+    if not api_key:
+        try:
+            api_key = getpass.getpass("KeyVault X-API-Key (won't echo): ")
+        except Exception:
+            api_key = None
+
+    if api_key and api_key.strip():
+        return _fetch_secret_json(keyvault_url, api_key.strip())
+
+    return None
+
+
+def _from_pretrained_kwargs_with_token(hf_token: str | None) -> dict:
+    if not hf_token:
+        return {}
+    # Newer HF stacks accept `token=...`; some older call sites still used `use_auth_token`.
+    return {"token": hf_token}
+
 # ========================= CONFIG =========================
 model_name = "meta-llama/Llama-3.2-1B-Instruct"
 dataset_name = "mlabonne/FineTome-100k"
@@ -22,6 +108,9 @@ max_steps = None               # set a number (e.g. 2000) if you want to train o
 
 output_dir = "./llama32-1b-finetuned-finetome"
 
+# Fetch an HF token at runtime (no persistence). Required for gated models.
+hf_token = get_hf_token()
+
 # ====================== LOAD MODEL (4-bit QLoRA) ======================
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -35,9 +124,14 @@ model = AutoModelForCausalLM.from_pretrained(
     quantization_config=bnb_config,
     device_map="auto",           # automatically puts layers on GPU
     trust_remote_code=True,
+    **_from_pretrained_kwargs_with_token(hf_token),
 )
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    **_from_pretrained_kwargs_with_token(hf_token),
+)
 tokenizer.pad_token = tokenizer.eos_token
 
 # Prepare for QLoRA
@@ -57,7 +151,11 @@ model = get_peft_model(model, lora_config)
 model.print_trainable_parameters()   # should show ~0.5-1% trainable params
 
 # ====================== LOAD & PREPARE DATASET ======================
-dataset = load_dataset(dataset_name, split="train")
+try:
+    dataset = load_dataset(dataset_name, split="train", token=hf_token) if hf_token else load_dataset(dataset_name, split="train")
+except TypeError:
+    # Fallback for older datasets APIs.
+    dataset = load_dataset(dataset_name, split="train", use_auth_token=hf_token) if hf_token else load_dataset(dataset_name, split="train")
 
 # Convert ShareGPT "conversations" → standard messages format
 def map_to_messages(example):
@@ -123,25 +221,13 @@ tokenizer.save_pretrained(output_dir)
 print(f"✅ Training finished! LoRA adapter saved to {output_dir}")
 
 
-#model merger from peft import AutoPeftModelForCausalLM
-
-model = AutoPeftModelForCausalLM.from_pretrained(
-    output_dir,
-    device_map="auto",
-    torch_dtype=torch.bfloat16
-)
-model = model.merge_and_unload()
-model.save_pretrained("llama32-1b-finetuned-merged")
-tokenizer.save_pretrained("llama32-1b-finetuned-merged")
-
-
-from peft import AutoPeftModelForCausalLM
-
+merged_dir = "llama32-1b-finetuned-merged"
 model = AutoPeftModelForCausalLM.from_pretrained(
     output_dir,
     device_map="auto",
-    torch_dtype=torch.bfloat16
+    torch_dtype=torch.bfloat16,
 )
 model = model.merge_and_unload()
-model.save_pretrained("llama32-1b-finetuned-merged")
-tokenizer.save_pretrained("llama32-1b-finetuned-merged")
+model.save_pretrained(merged_dir)
+tokenizer.save_pretrained(merged_dir)
+print(f"✅ Merged model saved to {merged_dir}")