Add finnetuing process configuration to model

app.py

@@ -10,7 +10,7 @@ import sys
 import torch
 
 
-from spanish_medica_llm import run_training, run_training_process
+from spanish_medica_llm import run_training, run_training_process, run_finnetuning_process
 
 import gradio as gr
 
@@ -45,41 +45,18 @@ def train_model(*inputs):
  if "IS_SHARED_UI" in os.environ:
  raise gr.Error("This Space only works in duplicated instances")
  
- # args_general = argparse.Namespace(
- # image_captions_filename = True,
- # train_text_encoder = True,
- # #stop_text_encoder_training = stptxt,
- # save_n_steps = 0,
- # #pretrained_model_name_or_path = model_to_load,
- # instance_data_dir="instance_images",
- # #class_data_dir=class_data_dir,
- # output_dir="output_model",
- # instance_prompt="",
- # seed=42,
- # resolution=512,
- # mixed_precision="fp16",
- # train_batch_size=1,
- # gradient_accumulation_steps=1,
- # use_8bit_adam=True,
- # learning_rate=2e-6,
- # lr_scheduler="polynomial",
- # lr_warmup_steps = 0,
- # #max_train_steps=Training_Steps, 
- # )
- # run_training(args_general)
- # torch.cuda.empty_cache()
- # #convert("output_model", "model.ckpt")
- # #shutil.rmtree('instance_images')
- # #shutil.make_archive("diffusers_model", 'zip', "output_model")
- # #with zipfile.ZipFile('diffusers_model.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:
- # # zipdir('output_model/', zipf)
- # torch.cuda.empty_cache()
- # return [gr.update(visible=True, value=["diffusers_model.zip"]), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)]
- run_training_process()
- 
+ run_training_process() 
  
  return f"Train Model Sucessful!!!"
 
+def finnetuning_model(*inputs):
+ if "IS_SHARED_UI" in os.environ:
+ raise gr.Error("This Space only works in duplicated instances")
+ 
+ run_finnetuning_process() 
+ 
+ return f"Finnetuning Model Sucessful!!!"
+
 def stop_model(*input):
  return f"Model with Gradio!"
 
@@ -93,6 +70,8 @@ with gr.Blocks() as demo:
  btn_response.click(fn=generate_model, inputs=inp, outputs=out)
  btn_train = gr.Button("Train Model")
  btn_train.click(fn=train_model, inputs=[], outputs=out)
+ btn_finnetuning = gr.Button("Finnetuning Model")
+ btn_finnetuning.click(fn=finnetuning_model, inputs=[], outputs=out)
  btn_evaluate = gr.Button("Evaluate Model")
  btn_evaluate.click(fn=evaluate_model, inputs=[], outputs=out)
  btn_stop = gr.Button("Stop Model")

spanish_medica_llm.py

@@ -331,6 +331,45 @@ MAX_TRAINING_STEPS = 2
 
 TOKEN_NAME = TOKEN_MISTRAL_NAME
 
+def get_chat_format(element):
+ """
+ Processes a single sample from the alpaca dataset to structure it for chatbot training.
+
+ This function transforms the dataset sample into a format suitable for training,
+ where each message is categorized by its role in the conversation (system, input, user, assistant).
+ It initializes the conversation with a system message, then conditionally adds an input message,
+ follows with the user's instruction, and finally, the assistant's output based on the provided inputs.
+
+ Parameters
+ ----------
+ sample : dict
+ A dictionary representing a single sample from the dataset. It must contain
+ keys corresponding to input and output components of the conversation.
+
+ Returns
+ -------
+ dict
+ A modified dictionary with a 'messages' key that contains a list of ordered messages,
+ each annotated with its role in the conversation.
+ """
+
+ prompt_template="""A partir del caso clínico que se expone a continuación, tu tarea es la siguiente.
+ Como médico experto, tu tarea es la de diagnosticar al paciente en base al caso clínico. Responde únicamente con el diagnóstico para el paciente de forma concisa.
+ Caso clínico: {caso_clinico}
+ """
+ # cómo usarlo con un LLM:
+
+ system_prompt = "Eres un experto en medicina que realiza diagnósticos en base a casos clínicos."
+
+ messages = [
+ {"role": "system", "content": system_prompt},
+ {"role": "user", "content": prompt_template.format(caso_clinico=element["raw_text"])},
+ {"role": "assistant", "content": element["topic"]},
+ ]
+
+ element["raw_text"] = messages
+ return element
+
 def loadSpanishTokenizer():
  """
  
@@ -379,12 +418,32 @@ def splitDatasetInTestValid(dataset):
  return (dataset['train'], eval_dataset, test_dataset)
 
 def loadSpanishDataset():
+ 
  spanishMedicaLllmDataset = load_dataset(SPANISH_MEDICA_LLM_DATASET, split="train")
  spanishMedicaLllmDataset = spanishMedicaLllmDataset.filter(lambda example: example["topic_type"] not in FILTER_CRITERIA)
  spanishMedicaLllmDataset = spanishMedicaLllmDataset.train_test_split(0.2, seed=203984)
  return spanishMedicaLllmDataset
 
+def loadSpanishDatasetFinnetuning():
+ 
+ spanishMedicaLllmDataset = load_dataset(SPANISH_MEDICA_LLM_DATASET, split="train")
+ spanishMedicaLllmDataset = spanishMedicaLllmDataset.filter(lambda example: example["topic_type"] in FILTER_CRITERIA)
+ return spanishMedicaLllmDataset
+
  ##See Jupyter Notebook for change CONTEXT_LENGTH size
+def applyChatInstructFormat(dataset, filterColumns = ['raw_text', 'topic']):
+ """
+ Apply instruccion chat_template
+ """
+ if dataset == None:
+ return dataset
+ else: 
+ dataset = dataset.remove_columns([col for col in dataset.features if col not in filterColumns])
+ return dataset.map(
+ get_chat_format,
+ batched=False,
+ num_proc=4
+ )
 
 def accelerateConfigModel():
  """
@@ -483,6 +542,26 @@ def modelLoraConfigBioMistral(model):
  model = accelerator.prepare_model(model)
  return (model)
  
+def getLoraConfiguration():
+ """
+ """
+ return LoraConfig(
+ r=8,
+ lora_alpha=16,
+ target_modules=[
+ "q_proj",
+ "k_proj",
+ "v_proj",
+ "o_proj",
+ "gate_proj",
+ "up_proj",
+ "down_proj",
+ "lm_head",
+ ],
+ bias="none",
+ lora_dropout=0.05, # Conventional
+ task_type="CAUSAL_LM",
+ )
 
 # A note on training. You can set the max_steps to be high initially, and examine at what step your
 # model's performance starts to degrade. There is where you'll find a sweet spot for how many steps 
@@ -541,10 +620,85 @@ def configAndRunTraining(basemodel, dataset, eval_dataset, tokenizer):
 
  trainer.push_to_hub()
  
+def configAndRunFineTuning(basemodel, dataset, eval_dataset, tokenizer):
+ if basemodel is None or dataset is None or tokenizer is None:
+ return None
+ else:
+ tokenizer.pad_token = tokenizer.eos_token
+
+ 
+ training_args = TrainingArguments(
+ output_dir=output_dir,
+ push_to_hub = True,
+ hub_private_repo = False,
+ hub_model_id = HUB_MODEL_ID,
+ warmup_steps = 5,
+ per_device_train_batch_size = MICRO_BATCH_SIZE,
+ per_device_eval_batch_size=1,
+ #gradient_checkpointing=True,
+ gradient_accumulation_steps = GRADIENT_ACCUMULATION_STEPS,
+ num_train_epochs = 1,
+ learning_rate = 2.5e-5, # Want about 10x smaller than the Mistral learning rate
+ logging_steps = 5,
+ optim="paged_adamw_8bit",
+ logging_dir="./logs", # Directory for storing logs
+ save_strategy = "steps", # Save the model checkpoint every logging step
+ save_steps = 50, # Save checkpoints every 50 steps
+ evaluation_strategy = "steps", # Evaluate the model every logging step
+ eval_steps = 50, # Evaluate and save checkpoints every 50 steps
+ do_eval = True, # Perform evaluation at the end of training
+ eval_steps=50,
+ save_total_limit=2,
+ remove_unused_columns = True,
+ report_to = None, # Comment this out if you don't want to use weights & baises
+ run_name=f"{run_name}-{datetime.now().strftime('%Y-%m-%d-%H-%M')}" , # Name of the W&B run (optional)
+ fp16=True, #Set for GPU T4 for more powerful GPU as G-100 or another change to false and bf16 parameter
+ bf16=False
+ )
+ 
+ trainer = SFTTrainer(
+ model=basemodel,
+ train_dataset = dataset,
+ eval_dataset = eval_dataset,
+ peft_config = getLoraConfiguration(),
+ dataset_text_field = "raw_text",
+ max_seq_length = 1024, #512
+ tokenizer = tokenizer,
+ args = training_args,
+ dataset_kwargs={
+ "add_special_tokens": False, # We template with special tokens
+ "append_concat_token": False, # No need to add additional separator token
+ },
+ packing=True
+ )
+ basemodel.config.use_cache = False # silence the warnings. Please re-enable for inference!
+ trainer.train()
 
 
+ trainer.push_to_hub()
+
 
 def run_training_process():
+ #Loggin to Huggin Face
+ login(token = os.environ.get('HG_FACE_TOKEN'))
+ os.environ['WANDB_DISABLED'] = 'true'
+ tokenizer = loadSpanishTokenizer()
+ medicalSpanishDataset = applyChatInstructFormat( loadSpanishDatasetFinnetuning())
+ medicalSpanishDataset = medicalSpanishDataset.train_test_split(0.2, seed=203984)
+
+ # train_dataset, eval_dataset, test_dataset = splitDatasetInTestValid(
+ # getTokenizedDataset( medicalSpanishDataset, tokenizer)
+ # )
+ 
+ 
+ train_dataset, eval_dataset, test_dataset = splitDatasetInTestValid( medicalSpanishDataset )
+
+ base_model = loadBaseModel(MISTRAL_BASE_MODEL_ID)
+ base_model = modelLoraConfigBioMistral(base_model)
+
+ configAndRunTraining(base_model,train_dataset, eval_dataset, tokenizer)
+
+def run_finnetuning_process():
  #Loggin to Huggin Face
  login(token = os.environ.get('HG_FACE_TOKEN'))
  os.environ['WANDB_DISABLED'] = 'true'
@@ -554,7 +708,7 @@ def run_training_process():
  getTokenizedDataset( medicalSpanishDataset, tokenizer)
  )
 
- base_model = loadBaseModel(MISTRAL_BASE_MODEL_ID)
- base_model = modelLoraConfigBioMistral(base_model)
+ base_model = loadBaseModel(HUB_MODEL_ID)
 
- configAndRunTraining(base_model,train_dataset, eval_dataset, tokenizer)
+ configAndRunFineTuning(base_model,train_dataset, eval_dataset, tokenizer)
+