Update README.md

README.md

@@ -22,6 +22,58 @@ This model combines the base `full_finetuned_llama3b` with LoRA fine-tuning on:
 from transformers import AutoTokenizer, AutoModelForCausalLM
 model = AutoModelForCausalLM.from_pretrained("your-hf-username/full_finetuned_llama3b")
 tokenizer = AutoTokenizer.from_pretrained("ai-factory/giant")
+# Load base model
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto",
+    trust_remote_code=True,
+    use_safetensors=True,
+    local_files_only=True
+)
+
+# Apply LoRA
+peft_config = LoraConfig(
+    task_type=TaskType.CAUSAL_LM,
+    r=8,
+    lora_alpha=32,
+    lora_dropout=0.05,
+    bias="none",
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]
+)
+model = get_peft_model(base_model, peft_config)
+model.eval()
+if torch.cuda.is_available():
+    model = model.cuda()
+
+# Load streaming datasets
+arxiv = load_dataset("ai-factory/red_pajama_subset_arxiv_subset", split="train", streaming=True)
+glaive = load_dataset("ai-factory/glaiveai-reasoning-v1-20m-chat", split="train", streaming=True)
+
+def tokenize(example):
+    return tokenizer(example["text"], truncation=True, max_length=4096)
+
+# Tokenize small samples
+tokenized_arxiv = map(tokenize, islice(arxiv, args.sample_size))
+tokenized_glaive = map(tokenize, islice(glaive, args.sample_size))
+
+# Run forward + backward pass (init LoRA weights)
+print("🔥 Training one step to initialize LoRA...")
+for i, sample in enumerate(tokenized_arxiv):
+    if not sample.get("input_ids"):
+        continue
+    ids = torch.tensor(sample["input_ids"]).unsqueeze(0).to(model.device)
+    labels = ids.clone()
+    loss = model(input_ids=ids, labels=labels).loss
+    loss.backward()
+    break
+
+# Merge LoRA and save
+print("🔁 Merging adapter into base model...")
+merged_model = model.merge_and_unload()
+merged_model.save_pretrained(SAVE_DIR, safe_serialization=True)
+tokenizer.save_pretrained(SAVE_DIR)
+print(f"✅ Merged model saved to {SAVE_DIR}")
 ```
 
 ## 👤 Authors