Refactor model and training scripts: remove unused imports and add FactorizedHead class for improved modularity

train.py

@@ -16,11 +16,6 @@ import math
 import sys
 from dataclasses import asdict, dataclass, field
 from typing import Literal, Sequence
-from pathlib import Path
-
-_PROJECT_ROOT = Path(__file__).resolve().parent
-if str(_PROJECT_ROOT) not in sys.path:
-    sys.path.insert(0, str(_PROJECT_ROOT))
 
 
 @dataclass