| |
| |
|
|
| import os |
| os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" |
| os.environ["CUDA_VISIBLE_DEVICES"] = "2" |
| from datasets import load_dataset |
| from unsloth import FastLanguageModel |
| from trl import SFTConfig, SFTTrainer |
|
|
| from unsloth.chat_templates import get_chat_template, train_on_responses_only |
|
|
| MODEL_NAME = "unsloth/Qwen3-8B" |
| DATA_PATH = "verified_combined_0-80.json" |
| TEST_DATA_PATH = "verified_combined_0-80_test.json" |
| MAX_SEQ_LENGTH = 4096 |
| FP16_SAVE_DIR = "/home/mshahidul/readctrl_model/full_model/classifier_model" |
| TEST_SPLIT_RATIO = 0.1 |
| SPLIT_SEED = 3407 |
|
|
| SYSTEM_PROMPT = ( |
| "You are an expert medical editor and Health Literacy specialist. " |
| "Classify the health literacy level of the provided text." |
| ) |
|
|
| USER_PROMPT = """Classify the health literacy level of the rewritten text. |
| |
| Labels: |
| - low_health_literacy: very simple, living-room language, minimal jargon. |
| - intermediate_health_literacy: standard public-friendly language, limited jargon. |
| - proficient_health_literacy: technical, clinical, or academic language. |
| |
| Input: |
| Full Source Text: |
| <<<FULLTEXT>>> |
| |
| Rewritten Text: |
| <<<DIFF_LABEL_TEXTS>>> |
| |
| Output: Return only one label string from the list above.""" |
|
|
|
|
| def build_messages(fulltext: str, diff_label_texts: str, label: str): |
| user_content = USER_PROMPT.replace("<<<FULLTEXT>>>", fulltext).replace( |
| "<<<DIFF_LABEL_TEXTS>>>", diff_label_texts |
| ) |
| return [ |
| {"role": "system", "content": SYSTEM_PROMPT}, |
| {"role": "user", "content": user_content}, |
| {"role": "assistant", "content": label}, |
| ] |
|
|
|
|
| def main(): |
| model, tokenizer = FastLanguageModel.from_pretrained( |
| model_name=MODEL_NAME, |
| max_seq_length=MAX_SEQ_LENGTH, |
| load_in_4bit=False, |
| load_in_8bit=False, |
| full_finetuning=False, |
| ) |
|
|
| model = FastLanguageModel.get_peft_model( |
| model, |
| r=32, |
| target_modules=[ |
| "q_proj", |
| "k_proj", |
| "v_proj", |
| "o_proj", |
| "gate_proj", |
| "up_proj", |
| "down_proj", |
| ], |
| lora_alpha=32, |
| lora_dropout=0, |
| bias="none", |
| use_gradient_checkpointing="unsloth", |
| random_state=3407, |
| use_rslora=False, |
| loftq_config=None, |
| ) |
|
|
| tokenizer = get_chat_template(tokenizer, chat_template="qwen3-instruct") |
| dataset = load_dataset("json", data_files=DATA_PATH, split="train") |
| split = dataset.train_test_split(test_size=TEST_SPLIT_RATIO, seed=SPLIT_SEED) |
| train_dataset = split["train"] |
| test_dataset = split["test"] |
| test_dataset.to_json(TEST_DATA_PATH) |
|
|
| def formatting_prompts_func(examples): |
| texts = [] |
| for fulltext, diff_label_texts, label in zip( |
| examples["fulltext"], |
| examples["diff_label_texts"], |
| examples["label"], |
| ): |
| messages = build_messages(fulltext, diff_label_texts, label) |
| text = tokenizer.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=False |
| ) |
| texts.append(text) |
| return {"text": texts} |
|
|
| train_dataset = train_dataset.map(formatting_prompts_func, batched=True) |
|
|
| trainer = SFTTrainer( |
| model=model, |
| processing_class=tokenizer, |
| train_dataset=train_dataset, |
| eval_dataset=None, |
| args=SFTConfig( |
| dataset_text_field="text", |
| per_device_train_batch_size=64, |
| gradient_accumulation_steps=16, |
| warmup_steps=5, |
| |
| num_train_epochs=1, |
| learning_rate=2e-4, |
| logging_steps=1, |
| optim="adamw_8bit", |
| weight_decay=0.001, |
| lr_scheduler_type="linear", |
| seed=3407, |
| report_to="none", |
| ), |
| ) |
|
|
| trainer = train_on_responses_only( |
| trainer, |
| instruction_part="<|im_start|>user\n", |
| response_part="<|im_start|>assistant\n", |
| ) |
|
|
| trainer.train() |
|
|
| os.makedirs(FP16_SAVE_DIR, exist_ok=True) |
| model.save_pretrained_merged( |
| FP16_SAVE_DIR, |
| tokenizer, |
| save_method="merged_16bit", |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
