IronWolfAI commited on
Commit
814b8e1
·
verified ·
1 Parent(s): 38ffa46

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -0
app.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from datasets import load_dataset
3
+ import torch
4
+ from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
5
+ from trl import SFTTrainer
6
+
7
+ # Load the model and tokenizer
8
+ model_name = "microsoft/phi-4-multimodal-instruct"
9
+ model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
10
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
11
+
12
+ # Load the dataset
13
+ dataset = load_dataset("openai/gsm8k", "main")["train"]
14
+
15
+ # Preprocess the dataset
16
+ def preprocess_function(examples):
17
+ return tokenizer(examples["question"], padding="max_length", truncation=True)
18
+
19
+ dataset = dataset.map(preprocess_function, batched=True)
20
+
21
+ # Define the training arguments
22
+ training_args = TrainingArguments(
23
+ output_dir="./results",
24
+ per_device_train_batch_size=4,
25
+ gradient_accumulation_steps=4,
26
+ learning_rate=2e-5,
27
+ num_train_epochs=1,
28
+ fp16=True,
29
+ logging_dir="./logs",
30
+ report_to="none",
31
+ )
32
+
33
+ # Create the SFT trainer
34
+ trainer = SFTTrainer(
35
+ model=model,
36
+ train_dataset=dataset,
37
+ args=training_args,
38
+ tokenizer=tokenizer,
39
+ )
40
+
41
+ # Train the model
42
+ trainer.train()
43
+
44
+ # Save the model
45
+ trainer.save_model("./results")