diff --git a/genai_stack/fine_tune/__init__.py b/genai_stack/fine_tune/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/genai_stack/fine_tune/base.py b/genai_stack/fine_tune/base.py new file mode 100644 index 00000000..a87f74ad --- /dev/null +++ b/genai_stack/fine_tune/base.py @@ -0,0 +1,40 @@ +from pydantic import BaseModel as PydanticBaseModel + + +class HyperParametersConfigModel(PydanticBaseModel): + pass + + +class BaseFineTuneConfigModel(PydanticBaseModel): + """ + Data Model for the configs + """ + + instruction: str = "text-summarization" + # hyperparameters: HyperParametersConfigModel + + +class BaseFineTune: + """ + Class which provides base structure for fine tuning models + """ + + config_class = BaseFineTuneConfigModel + + def __init__(self, **kwargs) -> None: + self.config = self.config_class(**kwargs) + + def load_model(self): + pass + + def load_dataset(self): + pass + + def train(self): + pass + + def tune(self): + pass + + def evaluate(self): + pass diff --git a/genai_stack/fine_tune/llama27b.py b/genai_stack/fine_tune/llama27b.py new file mode 100644 index 00000000..65aaed1e --- /dev/null +++ b/genai_stack/fine_tune/llama27b.py @@ -0,0 +1,48 @@ +from datasets import load_dataset +from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer +import torch + +from genai_stack.fine_tune.base import BaseFineTune, BaseFineTuneConfigModel +import ray + + +class LLama27bConfigModel(BaseFineTuneConfigModel): + model_name: str = "daryl149/llama-2-7b-chat-hf" + dataset: str + + +class LLama27bFineTune(BaseFineTune): + config_class = LLama27bConfigModel + + def load_model(self, model_name: str): + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + ) + torch.cuda.get_device_capability() + device_map = "auto" + + model = AutoModelForCausalLM.from_pretrained( + model_name, + quantization_config=bnb_config, + device_map=device_map, + ) + model.config.pretraining_tp = 1 + + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_fast=True) + tokenizer.pad_token = tokenizer + torch.cuda.empty_cache() + print("model", model) + print("tokenizer", tokenizer) + return model, tokenizer + + def load_dataset(self): + datasets = load_dataset(self.config.dataset) + ray_datasets = { + "train": ray.data.from_huggingface(datasets["train"]), + "validation": ray.data.from_huggingface(datasets["validation"]), + "test": ray.data.from_huggingface(datasets["test"]), + } + return ray_datasets