MasteredUltraInstinct commited on
Commit
eb0ee57
Β·
verified Β·
1 Parent(s): a84e0f7

Create train.py

Browse files
Files changed (1) hide show
  1. train.py +33 -0
train.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from model import get_model
2
+ from pix2tex.dataset.dataset import Im2LatexDataset
3
+ from pix2tex.trainer import Trainer
4
+ import os
5
+
6
+ os.makedirs('trained_model', exist_ok=True)
7
+
8
+ # Training parameters
9
+ config = {
10
+ "batch_size": 4,
11
+ "epochs": 1,
12
+ "max_seq_len": 150,
13
+ "warmup_steps": 10,
14
+ "lr": 1e-4,
15
+ "device": "cpu",
16
+ "save_dir": "trained_model",
17
+ "resume": False
18
+ }
19
+
20
+ # Dataset path
21
+ dataset = Im2LatexDataset(
22
+ data_root='handwritten_dataset',
23
+ transform=None,
24
+ max_length=config["max_seq_len"]
25
+ )
26
+
27
+ # Initialize model and trainer
28
+ model, tokenizer = get_model()
29
+ trainer = Trainer(model, tokenizer, config)
30
+
31
+ print("🧠 Starting training...")
32
+ trainer.train(dataset)
33
+ print("βœ… Training complete. Model saved to 'trained_model/'")