Spaces:
Running
Running
| from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments | |
| from peft import LoraConfig, get_peft_model | |
| from trl import SFTTrainer | |
| from datasets import load_dataset | |
| import torch | |
| # Load model and tokenizer | |
| model_name = "HuggingFaceTB/SmolLM3-3B" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| device_map="auto", | |
| torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, | |
| ) | |
| # Prepare PEFT config for efficient fine-tuning | |
| peft_config = LoraConfig( | |
| r=16, | |
| lora_alpha=32, | |
| target_modules=["q_proj", "v_proj", "k_proj", "o_proj"], | |
| lora_dropout=0.05, | |
| bias="none", | |
| task_type="CAUSAL_LM" | |
| ) | |
| model = get_peft_model(model, peft_config) | |
| # Load dataset (example: assume 'financial_data.jsonl' with {'text': 'query ||| response'} format) | |
| dataset = load_dataset("json", data_files="financial_data.jsonl", split="train") | |
| # Training arguments | |
| training_args = TrainingArguments( | |
| output_dir="./finetuned_smollm3", | |
| num_train_epochs=3, | |
| per_device_train_batch_size=4, | |
| gradient_accumulation_steps=4, | |
| learning_rate=2e-4, | |
| fp16=True if torch.cuda.is_available() else False, | |
| save_steps=500, | |
| logging_steps=100, | |
| optim="paged_adamw_8bit", | |
| weight_decay=0.01, | |
| warmup_steps=100, | |
| ) | |
| # Trainer | |
| trainer = SFTTrainer( | |
| model=model, | |
| args=training_args, | |
| train_dataset=dataset, | |
| peft_config=peft_config, | |
| dataset_text_field="text", # Adjust based on your dataset | |
| tokenizer=tokenizer, | |
| max_seq_length=512, | |
| ) | |
| trainer.train() | |
| # Save fine-tuned model | |
| trainer.model.save_pretrained("./finetuned_smollm3") | |
| tokenizer.save_pretrained("./finetuned_smollm3") |