finetune pipeline added

Author: Harikrishnan V

.gitignore

@@ -8,5 +8,6 @@ wheels/
 
 # Virtual environments
 .venv
-model_checkpoints/
+checkpoints/
 data/
+model_checkpoints/

check.py

@@ -0,0 +1,8 @@
+from transformers import PreTrainedTokenizerFast, AutoTokenizer
+from src.utils.paths import TOKENIZER_DIR
+from pathlib import Path
+
+
+tok = AutoTokenizer.from_pretrained("tokenizer_checkpoint")
+
+# tok = AutoTokenizer.from_pretrained(TOKENIZER_DIR)

configs/model.py

@@ -14,15 +14,21 @@ def __init__(
         ffn_multiple=128,
         tie_word_embeddings=True,
         rmsnorm_eps=1e-6,
+        rope_base=10000,
         **kwargs
     ):
         assert emb_dim % head_count == 0, "emb_dim must be divisible by head_count"
         self.vocab_size = vocab_size
         self.emb_dim = emb_dim
         self.num_blocks = num_blocks
         self.head_count = head_count
+        self.rmsnorm_eps = rmsnorm_eps
         self.seq_length = seq_length
         self.ffn_multiple = ffn_multiple
+        self.rope_base = rope_base
+        self.head_dim = self.emb_dim // self.head_count
+        self.architectures = ["BetterGPTModel"]
 
         # Initialize the Hugging Face base config
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+

src/finetuning/template.py configs/template.pysrc/finetuning/template.py renamed to configs/template.py

@@ -0,0 +1,90 @@
+import torch
+
+from trl import SFTTrainer
+from transformers import EarlyStoppingCallback
+from transformers import PreTrainedTokenizerFast,EarlyStoppingCallback
+from datasets import load_dataset
+
+from configs.ft_config import sft_config
+from src.utils.model_loader import load_model
+from src.utils.ft_tokenizer_helper import prepare_tokenizer
+from src.utils.prepare_alpaca_dataset import convert_alpaca_format
+from src.utils.paths import CHECKPOINT_DIR, TOKENIZER_DIR
+from src.utils.data_splitter import split_dataset
+from configs.tokenizer_config import TokenizerConfig
+from src.utils.logger import get_logger
+
+
+logger = get_logger(__name__)
+
+# model_path = CHECKPOINT_DIR
+# tokenizer_path = TOKENIZER_DIR
+model_path = "./checkpoints/model.pt"
+tokenizer_path = "./tokenizer_checkpoint"
+
+def main():
+    try:
+        logger.info("Starting fine-tuning process")
+        dataset = load_dataset("yahma/alpaca-cleaned", split="train")
+        logger.info("Alpaca Dataset loaded successfully")
+    except Exception as e:
+        logger.error(f"Error occurred while loading dataset or tokenizer: {e}")
+        raise
+    try:
+        dataset = dataset.map(convert_alpaca_format, batched=True)
+        logger.info("Alpaca Dataset converted to chat format successfully")
+    except Exception as e:
+        logger.error(f"Error occurred while converting dataset to chat format: {e}")
+        raise
+    try:
+        train_dataset, eval_dataset = split_dataset(dataset)
+        logger.info("Alpaca Dataset split into training and validation sets successfully")
+    except Exception as e:
+        logger.error(f"Error occurred while splitting dataset: {e}")
+        raise
+    try:
+        tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_path)
+        logger.info("Tokenizer loaded successfully")
+    except Exception as e:
+        logger.error(f"Error occurred while loading tokenizer: {e}")
+        raise
+    try:
+        tokenizer = prepare_tokenizer(tokenizer)
+        logger.info("Tokenizer prepared successfully")
+    except Exception as e:
+        logger.error(f"Error occurred while preparing tokenizer: {e}")
+        raise
+    try:
+        model = load_model(model_path, tokenizer)
+        logger.info("Model retied and loaded successfully")
+    except Exception as e:
+        logger.error(f"Error occurred while loading model: {e}")
+        raise
+    return model, train_dataset, eval_dataset, tokenizer
+
+
+if __name__ == "__main__":
+    model, train_dataset, eval_dataset, tokenizer = main()
+
+    model = load_model(model_path,tokenizer)
+
+    num_params = sum(p.numel() for p in model.parameters())
+    logger.info(f"{num_params/1e6:.2f}M parameters")
+
+    trainer = SFTTrainer(
+        model=model,
+        args=sft_config,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        processing_class=tokenizer,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)]
+    )
+
+    trainer.train()
+    logger.info("Fine-tuning completed successfully. Saving model and tokenizer...")
+
+    trainer.save_model(model_path)
+    tokenizer.save_pretrained(tokenizer_path)
+
+
+    

scripts/finetune_alpaca.py

@@ -4,7 +4,7 @@
 from torch.nn.functional import scaled_dot_product_attention
 
 from src.models.positional_embeddings import RoPESplitHalf
-from configs.model import ModelConfig
+from configs.model import BetterGPTConfig as ModelConfig
 from src.utils.logger import get_logger
 from src.utils.rope_helper import rotate_half, apply_rotary_pos_emb
 
@@ -20,9 +20,9 @@ class MHAttention(nn.Module):
 
     def __init__(
             self, 
-            emb_dim: int= ModelConfig.emb_dim, 
-            head_count: int= ModelConfig.head_count, 
-            head_dim: int= ModelConfig.head_dim
+            emb_dim: int, 
+            head_count: int, 
+            head_dim: int
             ):
         """
         Args:

src/finetuning/alpaca_ft.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from configs.model import ModelConfig
+from configs.model import BetterGPTConfig as ModelConfig
 
 class RMSNorm(nn.Module):
     """Root Mean Square Layer Normalization.
@@ -11,7 +11,7 @@ class RMSNorm(nn.Module):
     avoid precision loss with bfloat16 inputs.
     """
 
-    def __init__(self, emb_dim, eps: float = ModelConfig.rmsnorm_eps):
+    def __init__(self, emb_dim, eps: float = ModelConfig().rmsnorm_eps):
         """
         Args:
             emb_dim: Size of the last dimension to normalize over.

src/finetuning/finetune_model.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from configs.model import ModelConfig
+from configs.model import BetterGPTConfig as ModelConfig
 
 
 class RoPESplitHalf(nn.Module):
@@ -10,9 +10,7 @@ class RoPESplitHalf(nn.Module):
     half with the negated second half. compute sin/cos tables on the fly up to
     seq_len of each batch.
     """
-    def __init__(self, 
-                 head_dim: int = ModelConfig.head_dim,
-                 base: int = ModelConfig.rope_base):
+    def __init__(self, head_dim: int, base: float):
         """
         Args:
             head_dim: Per-head dimension; must be even.
@@ -40,4 +38,8 @@ def forward(self, x, seq_len):
             self._cos_cached, self._sin_cached = emb.cos(), emb.sin()
             self._cached_len = seq_len
         return self._cos_cached[:seq_len], self._sin_cached[:seq_len]
-
+    
+if __name__ == "__main__":
+    # config = ModelConfig()
+    print(f"RoPESplitHalf: head_dim={ModelConfig().head_dim}, rope_base={ModelConfig().rope_base}")
+    # RoPE = RoPESplitHalf(head_dim=config.head_dim, base=config.rope_base)