Update script files

src/predict.py

@@ -8,6 +8,7 @@ import pandas as pd
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from src.utils import read_binary_file
+from src.model import MalConv
 
 def predict_file(model_path, file_path, max_length=2_000_000):  # 2,000,000
     """
@@ -22,7 +23,11 @@ def predict_file(model_path, file_path, max_length=2_000_000):  # 2,000,000
         float: 예측 확률 (0에 가까우면 악성코드, 1에 가까우면 정상)
     """
     # 모델 로드
-    model = tf.keras.models.load_model(model_path)
+    model = MalConv(max_input_length=max_length)
+    # 모델의 가중치를 로드하기 전에 빌드
+    dummy_input = tf.zeros((1, max_length), dtype=tf.int32)
+    model(dummy_input)  # 모델 빌드
+    model.load_weights(model_path)
     
     # 파일 읽기
     byte_array = read_binary_file(file_path, max_length)
@@ -47,7 +52,11 @@ def predict_batch(model_path, csv_path, output_path=None, max_length=2**20):
     """
     # 모델 로드
     print("모델 로딩 중...")
-    model = tf.keras.models.load_model(model_path)
+    model = MalConv(max_input_length=max_length)
+    # 모델의 가중치를 로드하기 전에 빌드
+    dummy_input = tf.zeros((1, max_length), dtype=tf.int32)
+    model(dummy_input)  # 모델 빌드
+    model.load_weights(model_path)
     
     # CSV 파일 읽기
     df = pd.read_csv(csv_path)