Update app.py

app.py

@@ -73,6 +73,8 @@ def load_custom_model():
     model.eval()
     return model
 
+import importlib.util
+
 # 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()
@@ -96,40 +98,30 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
             print(f"Directory {model_repo_dir} does not exist.")
             return
 
-        # Step 3: Add the cloned repo to sys.path for imports
-        if model_repo_dir not in sys.path:
-            sys.path.append(model_repo_dir)
-
-        # Step 4: Debugging - Print sys.path to ensure the cloned repo is in the path
-        print(f"sys.path: {sys.path}")
-
-        # Ensure the 'LWM' directory is in the path
-        lwm_model_dir = os.path.join(os.getcwd(), 'LWM')
-        if lwm_model_dir not in sys.path:
-            sys.path.append(lwm_model_dir)
-
-        # Step 5: Verify if lwm_model.py exists in the directory
-        lwm_model_path = 'lwm_model.py'
+        # Step 3: Dynamically import lwm_model.py using importlib
+        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}")
             return f"Error: lwm_model.py not found at {lwm_model_path}"
-        else:
-            print(f"lwm_model.py found at {lwm_model_path}")
 
-        # Step 6: Import lwm_model after ensuring path is set
-        from lwm_model import LWM
+        # 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
         device = 'cpu'
         print(f"Loading the LWM model on {device}...")
-        model = LWM.from_pretrained(device=device)
+        model = lwm_model.LWM.from_pretrained(device=device)
 
-        # Step 7: Import tokenizer and load data
+        # Step 5: Import tokenizer and load data
         from input_preprocess import tokenizer
         with open(uploaded_file.name, 'rb') as f:
             manual_data = pickle.load(f)
 
         preprocessed_chs = tokenizer(manual_data=manual_data)
 
-        # Step 8: Perform inference
+        # 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)