From 79cb7eb0594cffd87eb12d24e92296c57426d29b Mon Sep 17 00:00:00 2001
From: MervinPraison <praison.ms@gmail.com>
Date: Wed, 5 Feb 2025 05:07:08 +0000
Subject: [PATCH] Update train.mdx with Wandb configuration and requirements

- Added detailed requirements for training, including Huggingface token and model details
- Included Wandb login and configuration instructions
- Updated documentation with environment variables for Wandb logging
---
 docs/train.mdx                     | 24 ++++++++++++++++++++++--
 praisonai/setup/setup_conda_env.sh |  2 +-
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/docs/train.mdx b/docs/train.mdx
index d755a715..6bbaf61a 100644
--- a/docs/train.mdx
+++ b/docs/train.mdx
@@ -28,8 +28,11 @@ praisonai train init
 
 ## Required
 
-1. HF_TOKEN
-2. Below config.yaml file in the same directory
+1. Huggingface token
+2. Base model to train on (e.g. unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit)
+3. Dataset to train on (e.g. yahma/alpaca-cleaned)
+4. Huggingface model name to upload to (e.g. mervinpraison/llama3.1-instruct) (Optional)
+5. Ollama model name to upload to (e.g. mervinpraison/llama3.1-instruct) (Optional)
 
 ## To upload to ollama.com (Linux)
 
@@ -119,4 +122,21 @@ quantization_method:
 
 ```bash
 praisonai train
+```
+
+## wandb
+
+```bash
+wandb login
+```
+
+<Note>
+Get the key from [here](https://wandb.ai/site/login)
+</Note>
+
+```bash
+export PRAISON_WANDB=True
+export WANDB_LOG_MODEL=checkpoint
+export WANDB_PROJECT=praisonai-test
+export PRAISON_WANDB_RUN_NAME=praisonai-train  
 ```
\ No newline at end of file
diff --git a/praisonai/setup/setup_conda_env.sh b/praisonai/setup/setup_conda_env.sh
index 0c379a49..0547a84e 100755
--- a/praisonai/setup/setup_conda_env.sh
+++ b/praisonai/setup/setup_conda_env.sh
@@ -100,6 +100,6 @@ $PIP_FULL_PATH install "unsloth[colab-new] @ git+https://github.com/unslothai/un
 $PIP_FULL_PATH install --no-deps "trl<0.9.0" peft accelerate bitsandbytes 
 $PIP_FULL_PATH install unsloth_zoo
 $PIP_FULL_PATH install cut_cross_entropy
-$PIP_FULL_PATH install sentencepiece protobuf datasets huggingface_hub hf_transfer
+$PIP_FULL_PATH install sentencepiece protobuf datasets huggingface_hub hf_transfer wandb
 
 echo "Setup completed successfully!"