some more cleanup

comparehf.py → scripts/comparehf.py

-from lm_arch import lm_class
+from lm_arch import lm_base
 from lm_arch.utils import * 
 import time
 
@@ -67,7 +67,7 @@ def timeit(func, r=1, n=5, quiet=False, function=None, do_tqdm=False, first=True
 
 
 with torch.no_grad():
-    based_model = lm_class.load_gpt_j().cuda().half().eval()
+    based_model = lm_base.load_gpt_j().cuda().half().eval()
     based_model = based_model.lm
     print("Loaded based model")
     hf_model = no_init(lambda: AutoModelForCausalLM.from_pretrained('/home/xuser/models/j6b_ckpt_14001')).cuda().half().eval()

train.py

@@ -6,7 +6,7 @@ import torch.optim as optim
 from pathlib import Path
 from lm_train import optimizer, utils
 from torch.utils import data
-from main import *
+from lm_arch import lm_base
 import yaml
 import sys
 from tqdm import tqdm