fix model path

NVIDIA · May 9, 2024 · 924734d · 924734d
1 parent 54fc3c9
commit 924734d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/nemo_curator/distributed_data_classification/domain_classifier_inference.py b/nemo_curator/distributed_data_classification/domain_classifier_inference.py
@@ -216,8 +216,9 @@ def main():
         columns_to_keep_list = df.columns.to_list()
         columns_to_keep_list.remove("sliced_text")
 
-        model_path = "/home/nfs/syurick/LLM_domain_classifier_inference/GoogleDebertaAgree_v3b_bce_maxlen512_bs64_best.pth"
-        model = DomainModel(Config, model_path=model_path, autocast=args.autocast)
+        model = DomainModel(
+            Config, model_path=args.model_file_name, autocast=args.autocast
+        )
         pipe = op.Sequential(
             op.Tokenizer(model, cols=["sliced_text"], tokenizer_type="sentencepiece"),
             op.Predictor(model, sorted_data_loader=True, batch_size=batch_size),