Update finetuned_model.py

finetuned_model.py

@@ -1,16 +1,21 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig
 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"
+# Load SmolLM-135M-Instruct model with 4-bit quantization
+model_name = "HuggingFaceTB/SmolLM-135M-Instruct"
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
+    quantization_config=quantization_config,
     device_map="auto",
-    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
 )
 
 # Prepare PEFT config for efficient fine-tuning
@@ -27,12 +32,12 @@ 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 arguments (adjusted for smaller model: larger batch size for speed)
 training_args = TrainingArguments(
-    output_dir="./finetuned_smollm3",
+    output_dir="./finetuned_smollm135m",
     num_train_epochs=3,
-    per_device_train_batch_size=4,
-    gradient_accumulation_steps=4,
+    per_device_train_batch_size=8,  # Increased for smaller model
+    gradient_accumulation_steps=2,
     learning_rate=2e-4,
     fp16=True if torch.cuda.is_available() else False,
     save_steps=500,
@@ -56,5 +61,5 @@ trainer = SFTTrainer(
 trainer.train()
 
 # Save fine-tuned model
-trainer.model.save_pretrained("./finetuned_smollm3")
-    tokenizer.save_pretrained("./finetuned_smollm3")
+trainer.model.save_pretrained("./finetuned_smollm135m")
+tokenizer.save_pretrained("./finetuned_smollm135m")