refactor(scripts): optimize finetune_model.py

Author: lzwjava

scripts/train_model.py renamed to scripts/finetune_model.py

@@ -2,13 +2,13 @@
 import glob
 import json
 from dotenv import load_dotenv
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling, LlamaTokenizerFast
 from datasets import Dataset, load_dataset
 import torch
 
 load_dotenv()
 
-MODEL_NAME = "gpt2"  # You can change this to a different model
+MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"  # Changed to the specified model
 OUTPUT_DIR = "trained_model"
 TRAIN_FILE = "train.jsonl"
 MAX_LENGTH = 512
@@ -47,7 +47,7 @@ def train_model(dataset, tokenizer):
         prediction_loss_only=True,
         remove_unused_columns=False,
     )
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True)
     data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
     trainer = Trainer(
         model=model,
@@ -61,7 +61,7 @@ def train_model(dataset, tokenizer):
 def main():
     posts_dir = "_posts"
     texts = create_training_data(posts_dir)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    tokenizer = LlamaTokenizerFast.from_pretrained(MODEL_NAME, trust_remote_code=True, use_fast=True)
     tokenizer.pad_token = tokenizer.eos_token
     dataset = prepare_dataset(texts, tokenizer)
     train_model(dataset, tokenizer)