diff --git a/diann_lib_processing.py b/diann_lib_processing.py
index bcc05cafa7197ddab5e65ff752186f8721e0ceae..90ef00290874e107ea563fa4627f2280dfa263c2 100644
--- a/diann_lib_processing.py
+++ b/diann_lib_processing.py
@@ -93,16 +93,16 @@ def predict(data_pred, model, output_path):
 
 
 if __name__ =='__main__':
-    df = load_lib('data/spectral_lib/first_lib.parquet')
-
-    plt.hist(df['RT'])
-    plt.savefig('test.png')
-
-    df_2 = pd.read_csv('data/data_prosit/data.csv')
-
-    plt.clf()
-    plt.hist(df_2['irt'])
-    plt.savefig('test2.png')
+    # df = load_lib('data/spectral_lib/first_lib.parquet')
+    #
+    # plt.hist(df['RT'])
+    # plt.savefig('test.png')
+    #
+    # df_2 = pd.read_csv('data/data_prosit/data.csv')
+    #
+    # plt.clf()
+    # plt.hist(df_2['irt'])
+    # plt.savefig('test2.png')
 
     # df_2 = extract_sequence(df).reset_index(drop=True)
     #
@@ -124,23 +124,20 @@ if __name__ =='__main__':
 
 
 
-    # args = load_args()
-    #
-    # model = ModelTransformer(encoder_ff=args.encoder_ff, decoder_rt_ff=args.decoder_rt_ff,
-    #                          n_head=args.n_head, encoder_num_layer=args.encoder_num_layer,
-    #                          decoder_rt_num_layer=args.decoder_rt_num_layer, drop_rate=args.drop_rate,
-    #                          embedding_dim=args.embedding_dim, acti=args.activation, norm=args.norm_first, seq_length=30)
-    #
-    # if torch.cuda.is_available():
-    #     model = model.cuda()
-    #
-    # model.load_state_dict(torch.load(args.model_weigh, weights_only=True))
-    #
-    # print(args.dataset_test)
-    # data_test = load_data(data_source='data/spectral_lib/data_uniprot_base.csv', batch_size=args.batch_size, length=30, mode=args.split_test,
-    #                       seq_col=args.seq_test)
-    #
-    # predict(data_test, model, args.output)
+    args = load_args()
+
+    model = ModelTransformer(encoder_ff=args.encoder_ff, decoder_rt_ff=args.decoder_rt_ff,
+                             n_head=args.n_head, encoder_num_layer=args.encoder_num_layer,
+                             decoder_rt_num_layer=args.decoder_rt_num_layer, drop_rate=args.drop_rate,
+                             embedding_dim=args.embedding_dim, acti=args.activation, norm=args.norm_first, seq_length=30)
+
+    if torch.cuda.is_available():
+        model = model.cuda()
+
+    model.load_state_dict(torch.load(args.model_weigh, weights_only=True))
+
+    print(args.dataset_test)
+    data_test = load_data(data_source='data/spectral_lib/data_uniprot_base.csv', batch_size=args.batch_size, length=30, mode=args.split_test,
+                          seq_col=args.seq_test)
 
-    plt.hist(df['RT'])
-    plt.savefig('test.png')
\ No newline at end of file
+    predict(data_test, model, args.output)