fullstuckdev/medication-ai-model

1from datasets import load_dataset2from transformers import (3    AutoTokenizer, 4    AutoModelForCausalLM,5    TrainingArguments,6    Trainer,7    DataCollatorForLanguageModeling8)9import torch10import os11 12model_output_path = "./model/medical_llama_3b"13os.makedirs(model_output_path, exist_ok=True)14 15model_name = "nvidia/Meta-Llama-3.2-3B-Instruct-ONNX-INT4"16tokenizer = AutoTokenizer.from_pretrained(model_name)17model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)18 19dataset = load_dataset("json", data_files="medical_dataset.json")20 21def preprocess_function(examples):22    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)23 24tokenized_dataset = dataset.map(25    preprocess_function,26    batched=True,27    remove_columns=dataset["train"].column_names28)29 30training_args = TrainingArguments(31    output_dir="./model/medical_llama_3b/checkpoints",  32    per_device_train_batch_size=4,33    gradient_accumulation_steps=4,34    num_train_epochs=3,35    learning_rate=2e-5,36    fp16=True,37    save_steps=500,38    logging_steps=100,39)40 41trainer = Trainer(42    model=model,43    args=training_args,44    train_dataset=tokenized_dataset["train"],45    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),46)47 48trainer.train()49 50model.save_pretrained(model_output_path)51tokenizer.save_pretrained(model_output_path)52 53print(f"Model and tokenizer saved to: {model_output_path}")