Update app.py

app.py

@@ -75,6 +75,13 @@ def load_custom_model():
 
 import importlib.util
 
+# Function to dynamically load a Python module from a given file path
+def load_module_from_path(module_name, file_path):
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
 # Function to process the uploaded .p file and perform inference using the custom model
 def process_p_file(uploaded_file, percentage_idx, complexity_idx):
     capture = PrintCapture()
@@ -89,7 +96,7 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
             print(f"Cloning model repository from {model_repo_url}...")
             subprocess.run(["git", "clone", model_repo_url, model_repo_dir], check=True)
 
-        # Step 2: Verify the repository was cloned
+        # Step 2: Verify the repository was cloned and change the working directory
         if os.path.exists(model_repo_dir):
             os.chdir(model_repo_dir)
             print(f"Changed working directory to {os.getcwd()}")
@@ -98,33 +105,43 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
             print(f"Directory {model_repo_dir} does not exist.")
             return
 
-        # Step 3: Dynamically import lwm_model.py using importlib
+        # Step 3: Dynamically load lwm_model.py, input_preprocess.py, and inference.py
         lwm_model_path = os.path.join(os.getcwd(), 'lwm_model.py')
-        if not os.path.exists(lwm_model_path):
-            print(f"Error: lwm_model.py not found at {lwm_model_path}")
+        input_preprocess_path = os.path.join(os.getcwd(), 'input_preprocess.py')
+        inference_path = os.path.join(os.getcwd(), 'inference.py')
+
+        # Load lwm_model
+        if os.path.exists(lwm_model_path):
+            lwm_model = load_module_from_path("lwm_model", lwm_model_path)
+        else:
             return f"Error: lwm_model.py not found at {lwm_model_path}"
 
-        # Use importlib to dynamically load lwm_model.py
-        spec = importlib.util.spec_from_file_location("lwm_model", lwm_model_path)
-        lwm_model = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(lwm_model)
-        
-        # Step 4: Load the model from LWM module
+        # Load input_preprocess
+        if os.path.exists(input_preprocess_path):
+            input_preprocess = load_module_from_path("input_preprocess", input_preprocess_path)
+        else:
+            return f"Error: input_preprocess.py not found at {input_preprocess_path}"
+
+        # Load inference
+        if os.path.exists(inference_path):
+            inference = load_module_from_path("inference", inference_path)
+        else:
+            return f"Error: inference.py not found at {inference_path}"
+
+        # Step 4: Load the model from lwm_model module
         device = 'cpu'
         print(f"Loading the LWM model on {device}...")
         model = lwm_model.LWM.from_pretrained(device=device)
 
-        # Step 5: Import tokenizer and load data
-        from input_preprocess import tokenizer
+        # Step 5: Tokenize the data using the tokenizer from input_preprocess
         with open(uploaded_file.name, 'rb') as f:
             manual_data = pickle.load(f)
 
-        preprocessed_chs = tokenizer(manual_data=manual_data)
+        preprocessed_chs = input_preprocess.tokenizer(manual_data=manual_data)
 
-        # Step 6: Perform inference
-        from inference import lwm_inference, create_raw_dataset
-        output_emb = lwm_inference(preprocessed_chs, 'channel_emb', model)
-        output_raw = create_raw_dataset(preprocessed_chs, device)
+        # Step 6: Perform inference using the functions from inference.py
+        output_emb = inference.lwm_inference(preprocessed_chs, 'channel_emb', model)
+        output_raw = inference.create_raw_dataset(preprocessed_chs, device)
 
         print(f"Output Embeddings Shape: {output_emb.shape}")
         print(f"Output Raw Shape: {output_raw.shape}")
@@ -137,6 +154,8 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
     finally:
         sys.stdout = sys.__stdout__  # Reset print statements
 
+
+
 # Function to handle logic based on whether a file is uploaded or not
 def los_nlos_classification(file, percentage_idx, complexity_idx):
     if file is not None: